{-# LANGUAGE BangPatterns #-}

module Statistics.Matrix.Fast (
    multiply,
    norm,
    multiplyV,
    transpose
    ) where

import Prelude hiding (exponent, map)
import Control.Monad.ST
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM

import Statistics.Matrix (row)
import Statistics.Matrix.Function
import Statistics.Matrix.Types
import Statistics.Matrix.Mutable  (unsafeNew,unsafeWrite,unsafeFreeze)

-- | Matrix-matrix multiplication in a more imperative fashion. Matrices must be of compatible
-- sizes (/note: not checked/). Faster but less accurate than Statistics.Matrix.multiply
multiply :: Matrix -> Matrix -> Matrix
multiply m1@(Matrix r1 _ _) m2@(Matrix _ c2 _) = runST $ do
  m3 <- unsafeNew r1 c2
  for 0 c2 $ \j -> do
    for 0 r1 $ \i -> do
      let 
        z = accum i m1 j m2 
      unsafeWrite m3 i j z
  unsafeFreeze m3

accum :: Int -> Matrix -> Int -> Matrix -> Double
accum ithrow (Matrix _ c1 v1) jthcol (Matrix _ c2 v2) = sub 0 0
  where sub !acc !ij | ij == c1 = acc
                     | otherwise = sub ( valRow*valCol + acc ) (ij+1)
                                   where 
                                    valRow = U.unsafeIndex v1 (ithrow*c1 + ij)
                                    valCol = U.unsafeIndex v2 (ij*c2+jthcol)

-- | Matrix-vector multiplication, with better performances but not as accurate as
-- Statistics.Matrix.multiplyV
multiplyV :: Matrix -> Vector -> Vector
multiplyV m v
  | cols m == c = U.generate (rows m) (U.sum . U.zipWith (*) v . row m)
  | otherwise   = error $ "matrix/vector unconformable " ++ show (cols m,c)
  where c = U.length v

-- | Norm of a vector. Faster but less accurate than Statistics.Matrix.norm
norm :: Vector -> Double
norm = sqrt . U.sum . U.map square

transpose :: Matrix -> Matrix
transpose (Matrix r0 c0 v0) 
  = Matrix c0 r0 $ runST $ do
    vec <- UM.unsafeNew (r0*c0)
    for 0 r0 $ \i -> do
      UM.unsafeWrite vec (i + i * c0) $ v0 `U.unsafeIndex` (i + i * c0)
      for (i+1) c0 $ \j -> do
        let tmp = v0 `U.unsafeIndex` (j + i * c0)
            tmp2 = v0 `U.unsafeIndex` (i + j * c0)
        UM.unsafeWrite vec (j + i * c0) tmp2
        UM.unsafeWrite vec (i + j * c0) tmp
    U.unsafeFreeze vec
