content: {-# LANGUAGE BangPatterns #-}
  
  -- |
  -- Module    : Statistics.Matrix.Fast.Algorithms
  -- Copyright : 2019 Magalame
  -- License   : BSD3
  --
  -- Useful matrix functions.
  
  module Statistics.Matrix.Fast.Algorithms
      (
        qr
      ) where
  
  import Control.Applicative ((<$>), (<*>))
  import Control.Monad.ST (ST, runST)
  import Prelude hiding (replicate)
  import Statistics.Matrix (Matrix (..),dimension, for)
  import qualified Statistics.Matrix.Mutable as M
  import qualified Data.Vector.Unboxed as U
  
  -- | /O(r*c)/ Compute the QR decomposition of a matrix.
  -- The result returned is the matrices (/q/,/r/).
  qr :: Matrix -> (Matrix, Matrix)
  qr mat = runST $ do
    let (m,n) = dimension mat
  
    r <- M.replicate n n 0
    a <- M.thaw mat
    for 0 n $ \j -> do
      cn <- M.immutably a $ \aa -> sqrt $ normCol j aa
      M.unsafeWrite r j j cn
      for 0 m $ \i -> M.unsafeModify a i j (/ cn)
      for (j+1) n $ \jj -> do
        p <- innerProduct a j jj
        M.unsafeWrite r j jj p
        for 0 m $ \i -> do
          aij <- M.unsafeRead a i j
          M.unsafeModify a i jj $ subtract (p * aij)
    (,) <$> M.unsafeFreeze a <*> M.unsafeFreeze r
  
  normCol :: Int -> Matrix -> Double
  normCol jthcol (Matrix r c v) = sub 0 0
    where sub !acc !ij | ij == r = acc
                       | otherwise = sub ( valCol*valCol + acc ) (ij+1)
                                     where
                                      valCol = U.unsafeIndex v (ij*c+jthcol)
  
  innerProduct :: M.MMatrix s -> Int -> Int -> ST s Double
  innerProduct mmat j k = M.immutably mmat $ \mat ->
    dotCol j mat k mat
  
  dotCol :: Int -> Matrix -> Int -> Matrix -> Double
  dotCol jthcol (Matrix r1 c1 v1) kthcol (Matrix _ c2 v2) = sub 0 0
    where sub !acc !ij | ij == r1 = acc
                       | otherwise = sub ( valColj*valColk + acc ) (ij+1)
                                     where
                                      valColk = U.unsafeIndex v2 (ij*c2+kthcol)
                                      valColj = U.unsafeIndex v1 (ij*c1+jthcol)