{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module Test.Matrix (testsVar) where

import qualified Test.Indexed as Indexed
import qualified Test.Generator as Gen
import qualified Test.Utility as Util
import Test.Generator
         ((<|*|>), (<|*.>), (<.*.>), (<***>), (<|=|>), (<|||>), (<===>))
import Test.Utility
         (equalArray, approx, approxArray, approxMatrix,
          maybeProperty, genOrder, Tagged(Tagged), TaggedGen)

import qualified Numeric.LAPACK.Matrix.Triangular as Triangular
import qualified Numeric.LAPACK.Matrix.Shape as MatrixShape
import qualified Numeric.LAPACK.Matrix.Square as Square
import qualified Numeric.LAPACK.Matrix as Matrix
import qualified Numeric.LAPACK.Vector as Vector
import Numeric.LAPACK.Matrix
         (General, ZeroInt, zeroInt, (#>), (<#>), (|||), (===))
import Numeric.LAPACK.Vector (Vector)
import Numeric.LAPACK.Scalar (RealOf, conjugate)

import qualified Numeric.Netlib.Class as Class

import qualified Data.Array.Comfort.Boxed as BoxedArray
import qualified Data.Array.Comfort.Storable as Array
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Storable (Array)
import Data.Array.Comfort.Shape ((:+:))

import Control.Applicative (liftA2, (<$>))

import Data.Traversable (forM)
import Data.Maybe.HT (toMaybe)
import Data.Tuple.HT (mapPair, swap)

import qualified Test.QuickCheck as QC


genArray ::
   (Shape.C shape, Class.Floating a) => shape -> QC.Gen (Array shape a)
genArray = Util.genArray 10


dotProduct ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Vector ZeroInt a, Vector ZeroInt a) -> Bool
dotProduct (x,y) =
   approx 1e-5
      (Vector.dot x y)
      (Matrix.toScalar $
       Matrix.singleRow MatrixShape.RowMajor x <#>
       Matrix.singleColumn MatrixShape.ColumnMajor y)

innerDot ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Vector ZeroInt a, Vector ZeroInt a) -> Bool
innerDot (x,y) =
   approx 1e-5 (Vector.inner x y) (Vector.dot (Vector.conjugate x) y)

tensorProductTranspose ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   MatrixShape.Order -> (Vector ZeroInt a, Vector ZeroInt a) -> Bool
tensorProductTranspose order (x,y) =
   approxArray
      (Matrix.transpose (Matrix.tensorProduct order x y))
      (Matrix.tensorProduct (MatrixShape.flipOrder order) y x)

outerTranspose ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   MatrixShape.Order -> (Vector ZeroInt a, Vector ZeroInt a) -> Bool
outerTranspose order (x,y) =
   approxArray
      (Matrix.transpose (Matrix.outer order x y))
      (Matrix.outer (MatrixShape.flipOrder order)
         (Vector.conjugate y) (Vector.conjugate x))

tensorProduct ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   MatrixShape.Order -> (Vector ZeroInt a, Vector ZeroInt a) -> Bool
tensorProduct order (x,y) =
   approxArray
      (Matrix.tensorProduct order x y)
      (Matrix.singleColumn order x <#> Matrix.singleRow order y)

tensorProductMul ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Triangular.Diagonal ZeroInt a,
    Matrix.General ZeroInt ZeroInt a,
    Triangular.Diagonal ZeroInt a) ->
   Bool
tensorProductMul (x,m,y) =
   let xmy = x <#> m <#> y
   in approxArray xmy
         (Vector.mul m
            (Matrix.tensorProduct (MatrixShape.fullOrder $ Array.shape xmy)
               (Triangular.takeDiagonal x) (Triangular.takeDiagonal y)))

outerTensorProduct ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   MatrixShape.Order -> (Vector ZeroInt a, Vector ZeroInt a) -> Bool
outerTensorProduct order (x,y) =
   approxArray
      (Matrix.outer order x y)
      (Matrix.tensorProduct order x $ Vector.conjugate y)

genScaledVectorPairs ::
   (Class.Floating a) =>
   Gen.Matrix a Int Int
      ((ZeroInt, ZeroInt), [(a, (Vector ZeroInt a, Vector ZeroInt a))])
genScaledVectorPairs =
   flip Gen.mapGen Gen.matrixDims $ \maxElem size@(height,width) ->
      fmap ((,) size) $
      QC.listOf $
         liftA2 (,) (Util.genElement maxElem) $
         liftA2 (,) (Util.genArray maxElem height) (Util.genArray maxElem width)

sumRank1 ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   MatrixShape.Order ->
   ((ZeroInt,ZeroInt), [(a, (Vector ZeroInt a, Vector ZeroInt a))]) -> Bool
sumRank1 order (size,xys) =
   approxArray
      (case order of
         MatrixShape.ColumnMajor -> Matrix.sumRank1 size xys
         MatrixShape.RowMajor ->
            Matrix.adjoint $
            Matrix.sumRank1 (swap size) $ map (mapPair (conjugate, swap)) xys)
      (foldl Vector.add
         (Vector.constant (uncurry (MatrixShape.general order) size) 0)
         (map (\(a,(x,y)) -> Matrix.outer order (Vector.scale a x) y) xys))


outerTrace ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   MatrixShape.Order -> (Vector ZeroInt a, Vector ZeroInt a) -> Bool
outerTrace order (x,y) =
   approx 1e-5
      (Vector.inner y x)
      (Square.trace $ Square.fromGeneral $ Matrix.outer order x y)

outerInner ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   MatrixShape.Order ->
   (Vector ZeroInt a, Vector ZeroInt a, Vector ZeroInt a) -> Bool
outerInner order (x,y,z) =
   approxArray (Matrix.outer order x y #> z) (Vector.scale (Vector.inner y z) x)


tensorTrace ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   MatrixShape.Order -> (Vector ZeroInt a, Vector ZeroInt a) -> Bool
tensorTrace order (x,y) =
   approx 1e-5 (Vector.dot y x)
      (Square.trace $ Square.fromGeneral $ Matrix.tensorProduct order x y)

tensorDot ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   MatrixShape.Order ->
   (Vector ZeroInt a, Vector ZeroInt a, Vector ZeroInt a) -> Bool
tensorDot order (x,y,z) =
   approxArray
      (Matrix.tensorProduct order x y #> z) (Vector.scale (Vector.dot y z) x)


genZeroColumns ::
   (Class.Floating a) => TaggedGen a (Matrix.Tall ZeroInt ZeroInt a)
genZeroColumns = Tagged $ do
   height <- zeroInt <$> QC.choose (0,5)
   order <- genOrder
   genArray (MatrixShape.tall order height (zeroInt 0))


reverseNoRows :: (Class.Floating a) => Matrix.Wide ZeroInt ZeroInt a -> Bool
reverseNoRows x =
   equalArray x $ Matrix.reverseRows x

reverseNoColumns :: (Class.Floating a) => Matrix.Tall ZeroInt ZeroInt a -> Bool
reverseNoColumns x =
   equalArray x $ Matrix.reverseColumns x



genMatrix2EqHeight ::
   (Class.Floating a) =>
   Gen.Matrix a Int (Int:+:Int)
      (General ZeroInt ZeroInt a, General ZeroInt ZeroInt a)
genMatrix2EqHeight = (,) <$> Gen.matrix <|||> Gen.matrix

genMatrix2EqWidth ::
   (Class.Floating a) =>
   Gen.Matrix a (Int:+:Int) Int
      (General ZeroInt ZeroInt a, General ZeroInt ZeroInt a)
genMatrix2EqWidth = (,) <$> Gen.matrix <===> Gen.matrix

reverseRows :: (Class.Floating a) => General ZeroInt ZeroInt a -> Bool
reverseRows x =
   equalArray x $ Matrix.reverseRows (Matrix.reverseRows x)

reverseColumns :: (Class.Floating a) => General ZeroInt ZeroInt a -> Bool
reverseColumns x =
   equalArray x $ Matrix.reverseColumns (Matrix.reverseColumns x)


-- cf. Vector.genSwapVector
genSwapVector ::
   (Class.Floating a) =>
   Gen.Vector a Int
      (Maybe (Int,Int), (General ZeroInt ZeroInt a, Vector ZeroInt a))
genSwapVector =
   flip Gen.mapGen ((,) <$> Gen.matrix <|*.> Gen.vector) $ \_maxElem (m,x) -> do
      let set = Shape.indices $ Array.shape x
      ij <-
         forM (toMaybe (not $ null set) set) $ \s ->
            liftA2 (,) (QC.elements s) (QC.elements s)
      return (ij,(m,x))

swapColumns ::
   (Class.Floating a) =>
   (Maybe (Int,Int), (General ZeroInt ZeroInt a, Vector ZeroInt a)) ->
   QC.Property
swapColumns (mij,(m,x)) =
   let mx = m#>x
   in seq (Array.shape mx) $
      maybeProperty $
      flip fmap mij $ \(i,j) ->
         equalArray mx (Matrix.swapColumns i j m #> Vector.swap i j x)

_swapColumns ::
   (Class.Floating a) =>
   ((Int,Int), (General ZeroInt ZeroInt a, Vector ZeroInt a)) -> Bool
_swapColumns ((i,j),(m,x)) =
   equalArray
      (m#>x)
      (Matrix.swapColumns i j m #> Vector.swap i j x)


zeroIntHeight ::
   (Shape.C height, Shape.C width) =>
   General height width a -> General ZeroInt width a
zeroIntHeight = Matrix.mapHeight (zeroInt . Shape.size)

zeroIntWidth ::
   (Shape.C height, Shape.C width) =>
   General height width a -> General height ZeroInt a
zeroIntWidth = Matrix.mapWidth (zeroInt . Shape.size)

reverseRowsStack ::
   (Class.Floating a) =>
   (General ZeroInt ZeroInt a, General ZeroInt ZeroInt a) -> Bool
reverseRowsStack (x,y) =
   equalArray
      (Matrix.reverseRows $ zeroIntHeight $ x===y)
      (zeroIntHeight $ Matrix.reverseRows y === Matrix.reverseRows x)

reverseColumnsStack ::
   (Class.Floating a) =>
   (General ZeroInt ZeroInt a, General ZeroInt ZeroInt a) -> Bool
reverseColumnsStack (x,y) =
   equalArray
      (Matrix.reverseColumns $ zeroIntWidth $ x|||y)
      (zeroIntWidth $ Matrix.reverseColumns y ||| Matrix.reverseColumns x)


data Cut = Take | Drop deriving (Show, Eq, Ord, Enum, Bounded)
data Slice = Row | Column deriving (Show, Eq, Ord, Enum, Bounded)

cut ::
   (Class.Floating a) =>
   Cut -> Slice -> Int ->
   General ZeroInt ZeroInt a -> General ZeroInt ZeroInt a
cut Take Row = Matrix.takeRows
cut Take Column = Matrix.takeColumns
cut Drop Row = Matrix.dropRows
cut Drop Column = Matrix.dropColumns

cutCommutative ::
   (Class.Floating a) =>
   ((Cut,Slice),(Int,Int)) -> General ZeroInt ZeroInt a -> Bool
cutCommutative (kind,(k,j)) x =
   let cutK = uncurry cut kind k
       cutJ = uncurry cut kind j
   in equalArray (cutK $ cutJ x) (cutJ $ cutK x)

cutRowColumnCommutative ::
   (Class.Floating a) =>
   ((Cut,Int),(Cut,Int)) -> General ZeroInt ZeroInt a -> Bool
cutRowColumnCommutative ((cutR,k),(cutC,j)) x =
   let cutRows = cut cutR Row k
       cutColumns = cut cutC Column j
   in equalArray (cutRows $ cutColumns x) (cutColumns $ cutRows x)


takeEqually ::
   (Class.Floating a) => Int -> General ZeroInt ZeroInt a -> Bool
takeEqually k x =
   equalArray
      (Matrix.takeEqually k x)
      (Matrix.takeRows k (Matrix.takeColumns k x))

dropEqually ::
   (Class.Floating a) => Int -> General ZeroInt ZeroInt a -> Bool
dropEqually k x =
   equalArray
      (Matrix.dropEqually k x)
      (Matrix.dropRows k (Matrix.dropColumns k x))


takeRowArray ::
   (Class.Floating a) => [Int] -> General ZeroInt ZeroInt a -> Bool
takeRowArray ixs x =
   equalArray
      (Matrix.adaptOrder x $
       Matrix.fromRows (Matrix.width x) $ map (Matrix.takeRow x) ixs)
      (Matrix.takeRowArray (BoxedArray.vectorFromList ixs) x)


stackSplitRows ::
   (Class.Floating a) => Int -> General ZeroInt ZeroInt a -> Bool
stackSplitRows k x =
   equalArray x
      (zeroIntHeight $ Matrix.takeRows k x === Matrix.dropRows k x)

stackSplitColumns ::
   (Class.Floating a) => Int -> General ZeroInt ZeroInt a -> Bool
stackSplitColumns k x =
   equalArray x
      (zeroIntWidth $ Matrix.takeColumns k x ||| Matrix.dropColumns k x)


takeStackRows, dropStackRows ::
   (Class.Floating a) =>
   (General ZeroInt ZeroInt a, General ZeroInt ZeroInt a) -> Bool
takeStackRows (x,y) =
   equalArray
      (Matrix.toRowMajor x)
      (Matrix.toRowMajor $ Matrix.takeRows (Shape.size $ Matrix.height x) $
       zeroIntHeight $ x===y)
dropStackRows (x,y) =
   equalArray
      (Matrix.toRowMajor y)
      (Matrix.toRowMajor $ Matrix.dropRows (Shape.size $ Matrix.height x) $
       zeroIntHeight $ x===y)

takeStackColumns, dropStackColumns ::
   (Class.Floating a) =>
   (General ZeroInt ZeroInt a, General ZeroInt ZeroInt a) -> Bool
takeStackColumns (x,y) =
   equalArray
      (Matrix.toRowMajor x)
      (Matrix.toRowMajor $ Matrix.takeColumns (Shape.size $ Matrix.width x) $
       zeroIntWidth $ x|||y)
dropStackColumns (x,y) =
   equalArray
      (Matrix.toRowMajor y)
      (Matrix.toRowMajor $ Matrix.dropColumns (Shape.size $ Matrix.width x) $
       zeroIntWidth $ x|||y)

stackRowsAssociative, stackColumnsAssociative ::
   (Class.Floating a) =>
   (General ZeroInt ZeroInt a,
    General ZeroInt ZeroInt a,
    General ZeroInt ZeroInt a) -> Bool
stackRowsAssociative (x,y,z) =
   equalArray
      (zeroIntHeight ((x===y)===z))
      (zeroIntHeight (x===(y===z)))
stackColumnsAssociative (x,y,z) =
   equalArray
      (zeroIntWidth ((x|||y)|||z))
      (zeroIntWidth (x|||(y|||z)))

stackRowsColumnsCommutative ::
   (Class.Floating a) =>
   ((General ZeroInt ZeroInt a, General ZeroInt ZeroInt a),
    (General ZeroInt ZeroInt a, General ZeroInt ZeroInt a)) -> Bool
stackRowsColumnsCommutative ((x,y),(z,w)) =
   equalArray
      (Matrix.toRowMajor $ (x|||y)===(z|||w))
      (Matrix.toRowMajor $ (x===z)|||(y===w))


forceOrder ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   MatrixShape.Order ->
   (General ZeroInt ZeroInt a, Vector ZeroInt a) ->
   Bool
forceOrder order (a,x) =
   let ao = Matrix.forceOrder order a
   in MatrixShape.fullOrder (Array.shape ao) == order
      &&
      approxArray (a #> x) (ao #> x)

addDistributive ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   ((General ZeroInt ZeroInt a, General ZeroInt ZeroInt a), Vector ZeroInt a) ->
   Bool
addDistributive ((a,b),x) =
   approxArray (Matrix.add a b #> x) (Vector.add (a#>x) (b#>x))

subDistributive ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   ((General ZeroInt ZeroInt a, General ZeroInt ZeroInt a), Vector ZeroInt a) ->
   Bool
subDistributive ((a,b),x) =
   approxArray (Matrix.sub a b #> x) (Vector.sub (a#>x) (b#>x))


multiplyDiagonalMatrix ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Triangular.Diagonal ZeroInt a, General ZeroInt ZeroInt a) -> Bool
multiplyDiagonalMatrix (x,y) =
   approxArray (x <#> y) (Triangular.toSquare x <#> y)

multiplyMatrixDiagonal ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (General ZeroInt ZeroInt a, Triangular.Diagonal ZeroInt a) -> Bool
multiplyMatrixDiagonal (x,y) =
   approxMatrix 1e-5 (x <#> y) (x <#> Triangular.toSquare y)



checkForAll ::
   (Show a, QC.Testable test) =>
   Gen.T tag dim a -> (a -> test) -> Tagged tag QC.Property
checkForAll gen = Util.checkForAll (Gen.run gen 10 5)

checkForAllExtra ::
   (Show a, Show b, QC.Testable test) =>
   QC.Gen a -> Gen.T tag dim b ->
   (a -> b -> test) -> Tagged tag QC.Property
checkForAllExtra = Gen.withExtra checkForAll


testsVar ::
   (Show a, Class.Floating a, Eq a, RealOf a ~ ar, Class.Real ar) =>
   [(String, Tagged a QC.Property)]
testsVar =
   ("index",
      checkForAll (Indexed.genMatrixIndex Gen.matrix) Indexed.unitDot) :
   ("dotProduct",
      checkForAll ((,) <$> Gen.vector <.*.> Gen.vector) dotProduct) :
   ("innerDot",
      checkForAll ((,) <$> Gen.vector <.*.> Gen.vector) innerDot) :
   ("tensorProductTranspose",
      checkForAllExtra genOrder
         ((,) <$> Gen.vector <***> Gen.vector) tensorProductTranspose) :
   ("outerTranspose",
      checkForAllExtra genOrder
         ((,) <$> Gen.vector <***> Gen.vector) outerTranspose) :
   ("tensorProduct",
      checkForAllExtra genOrder
         ((,) <$> Gen.vector <***> Gen.vector) tensorProduct) :
   ("tensorProductMul",
      checkForAll ((,,) <$> Gen.diagonal <|*|> Gen.matrix <|*|> Gen.diagonal)
         tensorProductMul) :
   ("outerTensorProduct",
      checkForAllExtra genOrder
         ((,) <$> Gen.vector <***> Gen.vector) outerTensorProduct) :
   ("sumRank1",
      checkForAllExtra genOrder genScaledVectorPairs sumRank1) :

   ("outerTrace",
      checkForAllExtra genOrder
         ((,) <$> Gen.vector <.*.> Gen.vector) outerTrace) :
   ("outerInner",
      checkForAllExtra genOrder
         ((,,) <$> Gen.vector <***> Gen.vector <|*.> Gen.vector) outerInner) :
   ("tensorTrace",
      checkForAllExtra genOrder
         ((,) <$> Gen.vector <.*.> Gen.vector) tensorTrace) :
   ("tensorDot",
      checkForAllExtra genOrder
         ((,,) <$> Gen.vector <***> Gen.vector <|*.> Gen.vector) tensorDot) :

   ("reverseNoRows",
      Util.checkForAllPlain
         (fmap Matrix.transpose <$> genZeroColumns) reverseNoRows) :
   ("reverseNoColumns",
      Util.checkForAllPlain genZeroColumns reverseNoColumns) :
   ("reverseRows",
      checkForAll Gen.matrix reverseRows) :
   ("reverseColumns",
      checkForAll Gen.matrix reverseColumns) :
   ("reverseRowsStack",
      checkForAll genMatrix2EqWidth reverseRowsStack) :
   ("reverseColumnsStack",
      checkForAll genMatrix2EqHeight reverseColumnsStack) :
   ("cutCommutative",
      checkForAllExtra
         (liftA2 (,)
            (liftA2 (,) QC.arbitraryBoundedEnum QC.arbitraryBoundedEnum)
            (liftA2 (,) (QC.choose (0,5)) (QC.choose (0,5))))
         Gen.matrix cutCommutative) :
   ("cutRowColumnCommutative",
      checkForAllExtra
         (liftA2 (,)
            (liftA2 (,) QC.arbitraryBoundedEnum (QC.choose (0,5)))
            (liftA2 (,) QC.arbitraryBoundedEnum (QC.choose (0,5))))
         Gen.matrix cutRowColumnCommutative) :
   ("takeEqually",
      checkForAllExtra (QC.choose (0,5)) Gen.matrix takeEqually) :
   ("dropEqually",
      checkForAllExtra (QC.choose (0,5)) Gen.matrix dropEqually) :
   ("takeRowArray",
      checkForAll
         (Gen.mapGen
            (\_maxElem x -> do
               let set = Shape.indices $ Matrix.height x
               ixs <-
                  if null set then return [] else QC.listOf $ QC.elements set
               return (ixs,x))
            Gen.matrix)
         (uncurry takeRowArray)) :
   ("swapColumns",
      checkForAll genSwapVector swapColumns) :
   ("stackSplitRows",
      checkForAllExtra (QC.choose (0,5)) Gen.matrix stackSplitRows) :
   ("stackSplitColumns",
      checkForAllExtra (QC.choose (0,5)) Gen.matrix stackSplitColumns) :
   ("takeStackRows",
      checkForAll genMatrix2EqWidth takeStackRows) :
   ("dropStackRows",
      checkForAll genMatrix2EqWidth dropStackRows) :
   ("takeStackColumns",
      checkForAll genMatrix2EqHeight takeStackColumns) :
   ("dropStackColumns",
      checkForAll genMatrix2EqHeight dropStackColumns) :
   ("stackRowsAssociative",
      checkForAll
         ((,,) <$> Gen.matrix <===> Gen.matrix <===> Gen.matrix)
         stackRowsAssociative) :
   ("stackColumnsAssociative",
      checkForAll
         ((,,) <$> Gen.matrix <|||> Gen.matrix <|||> Gen.matrix)
         stackColumnsAssociative) :
   ("stackRowsColumnsCommutative",
      checkForAll
         ((,) <$> genMatrix2EqHeight <===> genMatrix2EqHeight)
         stackRowsColumnsCommutative) :

   ("forceOrder",
      checkForAllExtra genOrder
         ((,) <$> Gen.matrix <|*.> Gen.vector) forceOrder) :
   ("addDistributive",
      checkForAll
         ((,) <$>
            ((,) <$> Gen.matrix <|=|> Gen.matrix)
            <|*.>
            Gen.vector)
         addDistributive) :
   ("subDistributive",
      checkForAll
         ((,) <$>
            ((,) <$> Gen.matrix <|=|> Gen.matrix)
            <|*.>
            Gen.vector)
         subDistributive) :

   ("multiplyDiagonalMatrix",
      checkForAll
         ((,) <$> Gen.diagonal <|*|> Gen.matrix) multiplyDiagonalMatrix) :
   ("multiplyMatrixDiagonal",
      checkForAll
         ((,) <$> Gen.matrix <|*|> Gen.diagonal) multiplyMatrixDiagonal) :
   []