```{-# OPTIONS_GHC -fglasgow-exts #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Data.Matrix.Dense.Operations
-- Maintainer : Patrick Perry <patperry@stanford.edu>
-- Stability  : experimental
--

module Data.Matrix.Dense.Operations (
-- * Copy and swap
copyMatrix,
swapMatrices,

-- * Matrix Arithmetic
-- ** Pure
shift,
scale,
invScale,
plus,
minus,
times,
divide,

-- ** Impure
getShifted,
getScaled,
getInvScaled,
getSum,
getDiff,
getProduct,
getRatio,

-- * In-place operations
doConj,
shiftBy,
scaleBy,
invScaleBy,
axpy,
(+=),
(-=),
(*=),
(//=),

-- * BLAS operations
gemv,
gemm,

-- * Unsafe operations
unsafeCopyMatrix,
unsafeSwapMatrices,

) where

import Data.Maybe ( fromJust )
import System.IO.Unsafe
import Unsafe.Coerce

import BLAS.Internal ( checkMatMatOp )
import Data.Matrix.Dense.Internal
import Data.Vector.Dense.Internal hiding ( unsafeWithElemPtr, unsafeThaw,
unsafeFreeze )
import qualified Data.Vector.Dense.Operations as V
import qualified Data.Vector.Dense.Internal as V

import BLAS.C ( CBLASTrans, colMajor, noTrans, conjTrans )
import qualified BLAS.C as BLAS
import BLAS.Elem ( BLAS1, BLAS2, BLAS3 )
import qualified BLAS.Elem as E

infixl 7 `scale`, `invScale`
infixl 6 `shift`
infixl 1 +=, -=, *=, //=

-- | @copy dst src@ copies the elements from the second argument to the first.
copyMatrix :: (BLAS1 e) => IOMatrix (m,n) e -> DMatrix t (m,n) e -> IO ()
copyMatrix a b =
checkMatMatOp "copyMatrix" (shape a) (shape b) \$ unsafeCopyMatrix a b

unsafeCopyMatrix :: (BLAS1 e) => IOMatrix (m,n) e -> DMatrix t (m,n) e -> IO ()
unsafeCopyMatrix = liftV2 (V.unsafeCopyVector)

-- | @swap a b@ exchanges the elements stored in two matrices.
swapMatrices :: (BLAS1 e) => IOMatrix (m,n) e -> IOMatrix (m,n) e -> IO ()
swapMatrices a b =
checkMatMatOp "swapMatrices" (shape a) (shape b) \$ unsafeSwapMatrices a b

unsafeSwapMatrices :: (BLAS1 e) => IOMatrix (m,n) e -> IOMatrix (m,n) e -> IO ()
unsafeSwapMatrices = liftV2 (V.unsafeSwapVectors)

-- | Form a new matrix by adding a value to every element in a matrix.
getShifted :: (BLAS1 e) => e -> DMatrix t (m,n) e -> IO (DMatrix r (m,n) e)
getShifted k = unaryOp (shiftBy k)

-- | Form a new matrix by multiplying every element by a value.
getScaled :: (BLAS1 e) => e -> DMatrix t (m,n) e -> IO (DMatrix r (m,n) e)
getScaled k = unaryOp (scaleBy k)

-- | Form a new matrix by dividing every element by a value.
getInvScaled :: (BLAS1 e) => e -> DMatrix t (m,n) e -> IO (DMatrix r (m,n) e)
getInvScaled k = unaryOp (invScaleBy k)

-- | Create a new matrix by taking the elementwise sum of two matrices.
getSum :: (BLAS1 e) => e -> DMatrix s (m,n) e -> e -> DMatrix t (m,n) e -> IO (DMatrix r (m,n) e)
getSum alpha a beta b =
checkMatMatOp "getSum" (shape a) (shape b) \$ unsafeGetSum alpha a beta b

unsafeGetSum :: (BLAS1 e) => e -> DMatrix s (m,n) e -> e -> DMatrix t (m,n) e -> IO (DMatrix r (m,n) e)
unsafeGetSum alpha a beta b
| isHerm a = do
s <- unsafeGetSum (E.conj alpha) (herm a) (E.conj beta) (herm b)
return (herm s)
| otherwise = do
s <- getScaled alpha a
axpy beta b (unsafeThaw s)
return (unsafeCoerce s)

-- | Create a new matrix by taking the elementwise difference of two matrices.
getDiff :: (BLAS1 e) => DMatrix s (m,n) e -> DMatrix t (m,n) e -> IO (DMatrix r (m,n) e)
getDiff a b = checkMatMatOp "minus" (shape a) (shape b) \$ unsafeGetSum 1 a (-1) b

-- | Create a new matrix by taking the elementwise product of two matrices.
getProduct :: (BLAS2 e) => DMatrix s (m,n) e -> DMatrix t (m,n) e -> IO (DMatrix r (m,n) e)
getProduct = binaryOp "times" (*=)

-- | Create a new matrix by taking the elementwise ratio of two matrices.
getRatio :: (BLAS2 e) => DMatrix s (m,n) e -> DMatrix t (m,n) e -> IO (DMatrix r (m,n) e)
getRatio = binaryOp "getRatio" (//=)

-- | Conjugate every element in a matrix.
doConj  :: (BLAS1 e) => IOMatrix (m,n) e -> IO ()
doConj a | isHerm a  = doConj (herm a)
| otherwise = liftV (\x -> V.doConj x >> return ()) a

-- | Scale every element in a matrix by the given value.
scaleBy :: (BLAS1 e) => e -> IOMatrix (m,n) e -> IO ()
scaleBy k = liftV (\x -> V.scaleBy k x >> return ())

-- | Scale every element by the given value.
shiftBy :: (BLAS1 e) => e -> IOMatrix (m,n) e -> IO ()
shiftBy k = liftV (V.shiftBy k)

-- | Divide every element by the given value.
invScaleBy :: (BLAS1 e) => e -> IOMatrix (m,n) e -> IO ()
invScaleBy k = liftV (V.invScaleBy k)

axpy :: (BLAS1 e) => e -> DMatrix t (m,n) e -> IOMatrix (m,n) e -> IO ()
axpy alpha = liftV2 (V.axpy alpha)

(+=) :: (BLAS1 e) => IOMatrix (m,n) e -> DMatrix t (m,n) e -> IO ()
(+=) = liftV2 (V.+=)

-- | In-place elementwise subtract.
(-=) :: (BLAS1 e) => IOMatrix (m,n) e -> DMatrix t (m,n) e -> IO ()
(-=) = liftV2 (V.-=)

-- | In-place elementwise product.
(*=) :: (BLAS2 e) => IOMatrix (m,n) e -> DMatrix t (m,n) e -> IO ()
(*=) = liftV2 (V.*=)

-- | In-place elementwise divide.
(//=) :: (BLAS2 e) => IOMatrix (m,n) e -> DMatrix t (m,n) e -> IO ()
(//=) = liftV2 (V.//=)

blasTransOf :: DMatrix t (m,n) e -> CBLASTrans
blasTransOf a =
case (isHerm a) of
False -> noTrans
True  -> conjTrans

flipShape :: (Int,Int) -> (Int,Int)
flipShape (m,n) = (n,m)

-- | @gemv alpha a x beta y@ replaces @y := alpha a * x + beta y@.
gemv :: (BLAS3 e) => e -> DMatrix s (m,n) e -> DVector t n e -> e -> IOVector m e -> IO ()
gemv alpha a x beta y
| numRows a == 0 || numCols a == 0 =
return ()
| isConj x =
let b = fromJust \$ maybeFromCol x
in case maybeFromCol y of
Just c  -> gemm alpha a b beta c
Nothing -> do
V.doConj y
gemm alpha a b beta (fromJust \$ maybeFromCol \$ V.conj y)
V.doConj y
| isConj y =
let c = fromJust \$ maybeFromCol y
in case maybeFromCol x of
Just b  -> gemm alpha a b beta c
Nothing -> do
x' <- V.newCopy (V.unsafeThaw x)
V.doConj x'
gemm alpha a (fromJust \$ maybeFromCol \$ V.conj x') beta c
| otherwise =
let order  = colMajor
transA = blasTransOf a
(m,n)  = case (isHerm a) of
False -> shape a
True  -> (flipShape . shape) a
ldA    = ldaOf a
incX   = V.strideOf x
incY   = V.strideOf y
in unsafeWithElemPtr a (0,0) \$ \pA ->
V.unsafeWithElemPtr x 0 \$ \pX ->
V.unsafeWithElemPtr y 0 \$ \pY -> do
BLAS.gemv order transA m n alpha pA ldA pX incX beta pY incY

-- | @gemm alpha a b beta c@ replaces @c := alpha a * b + beta c@.
gemm :: (BLAS3 e) => e -> DMatrix s (m,k) e -> DMatrix t (k,n) e -> e -> IOMatrix (m,n) e -> IO ()
gemm alpha a b beta c
| numRows a == 0 || numCols a == 0 || numCols b == 0 = return ()
| isHerm c = gemm (E.conj alpha) (herm b) (herm a) (E.conj beta) (herm c)
| otherwise =
let order  = colMajor
transA = blasTransOf a
transB = blasTransOf b
(m,n)  = shape c
k      = numCols a
ldA    = ldaOf a
ldB    = ldaOf b
ldC    = ldaOf c
in unsafeWithElemPtr a (0,0) \$ \pA ->
unsafeWithElemPtr b (0,0) \$ \pB ->
unsafeWithElemPtr c (0,0) \$ \pC ->
BLAS.gemm order transA transB m n k alpha pA ldA pB ldB beta pC ldC

unaryOp :: (BLAS1 e) => (IOMatrix (m,n) e -> IO ())
-> DMatrix t (m,n) e -> IO (DMatrix r (m,n) e)
unaryOp f a = do
a' <- newCopy a
f (unsafeThaw a')
return (unsafeCoerce a')

binaryOp :: (BLAS1 e) => String -> (IOMatrix (m,n) e -> DMatrix t (m,n) e -> IO ())
-> DMatrix s (m,n) e -> DMatrix t (m,n) e -> IO (DMatrix r (m,n) e)
binaryOp name f a b =
checkMatMatOp name (shape a) (shape b) \$ do
a' <- newCopy a
f (unsafeThaw a') b
return (unsafeCoerce a')

-- | Create a new matrix by scaling another matrix by the given value.
scale :: (BLAS1 e) => e -> Matrix (m,n) e -> Matrix (m,n) e
scale k a = unsafePerformIO \$ getScaled k a
{-# NOINLINE scale #-}

-- | Form a new matrix by adding a value to every element in a matrix.
shift :: (BLAS1 e) => e -> Matrix (m,n) e -> Matrix (m,n) e
shift k a = unsafePerformIO \$ getShifted k a
{-# NOINLINE shift #-}

-- | Form a new matrix by dividing every element by a value.
invScale :: (BLAS1 e) => e -> Matrix (m,n) e -> Matrix (m,n) e
invScale k a = unsafePerformIO \$ getInvScaled k a
{-# NOINLINE invScale #-}

-- | Create a new matrix by taking the elementwise sum of two matrices.
add :: (BLAS1 e) => e -> Matrix (m,n) e -> e -> Matrix (m,n) e -> Matrix (m,n) e
add alpha a beta b = unsafePerformIO \$ getSum alpha a beta b

-- | Create a new matrix by taking the elementwise sum of two matrices.
plus :: (BLAS1 e) => Matrix (m,n) e -> Matrix (m,n) e -> Matrix (m,n) e
plus a b = add 1 a 1 b

-- | Create a new matrix by taking the elementwise difference of two matrices.
minus :: (BLAS1 e) => Matrix (m,n) e -> Matrix (m,n) e -> Matrix (m,n) e
minus a b = unsafePerformIO \$ getDiff a b
{-# NOINLINE minus #-}

-- | Create a new matrix by taking the elementwise product of two matrices.
times :: (BLAS2 e) => Matrix (m,n) e -> Matrix (m,n) e -> Matrix (m,n) e
times a b = unsafePerformIO \$ getProduct a b
{-# NOINLINE times #-}

-- | Create a new matrix by taking the elementwise ratio of two matrices.
divide :: (BLAS2 e) => Matrix (m,n) e -> Matrix (m,n) e -> Matrix (m,n) e
divide a b = unsafePerformIO \$ getRatio a b
{-# NOINLINE divide #-}

{-# RULES
"scale.plus/add"   forall k l x y. plus (scale k x) (scale l y) = add k x l y
"scale1.plus/add"  forall k x y.   plus (scale k x) y = add k x 1 y
"scale2.plus/add"  forall k x y.   plus x (scale k y) = add 1 x k y

"scale.minus/add"  forall k l x y. minus (scale k x) (scale l y) = add k x (-l) y
"scale1.minus/add" forall k x y.   minus (scale k x) y = add k x (-1) y
"scale2.minus/add" forall k x y.   minus x (scale k y) = add 1 x (-k) y
#-}

```