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

import qualified Test.Divide as Divide
import qualified Test.Multiply as Multiply
import qualified Test.Generic as Generic
import qualified Test.Generator as Gen
import qualified Test.Logic as Logic
import qualified Test.Utility as Util
import Test.Generator ((<|||>), (<===>))
import Test.Utility (Tagged, prefix)

import qualified Numeric.LAPACK.Matrix.Block as Block
import qualified Numeric.LAPACK.Matrix.Square as Square
import qualified Numeric.LAPACK.Matrix.Array as ArrMatrix
import qualified Numeric.LAPACK.Matrix.Shape as MatrixShape
import qualified Numeric.LAPACK.Matrix.Layout as Layout
import qualified Numeric.LAPACK.Matrix.Extent as Extent
import qualified Numeric.LAPACK.Matrix as Matrix
import Numeric.LAPACK.Matrix.Shape (Filled)
import Numeric.LAPACK.Matrix.Extent (Big)
import Numeric.LAPACK.Matrix (ShapeInt)
import Numeric.LAPACK.Scalar (RealOf)

import qualified Numeric.Netlib.Class as Class

import Data.Array.Comfort.Shape ((::+))

import Control.Applicative ((<$>))

import qualified Test.QuickCheck as QC


type TypeFull = ArrMatrix.Array Layout.Unpacked MatrixShape.Arbitrary

genIdentity ::
   (Logic.Dim sh, Class.Floating a) => Gen.Square sh a (Matrix.Square sh a)
genIdentity =
   Gen.fromBase (return . Matrix.identityFromShape <$> Gen.squareShape)


type Diagonal =
      Block.Diagonal TypeFull () () TypeFull () ()
         Filled Filled ShapeInt ShapeInt

diagonalTransposeUnpack :: (Class.Floating a) => Diagonal a -> Bool
diagonalTransposeUnpack a =
   Util.equalArray
      (Matrix.toFull (Matrix.transpose a))
      (Matrix.transpose (Matrix.toFull a))


genDiagonal ::
   (Class.Floating a) => Gen.Square (ShapeInt::+ShapeInt) a (Diagonal a)
genDiagonal =
   uncurry Block.Diagonal <$> Gen.stackDiagonal Gen.square Gen.square

genIdentityDiagonal ::
   (Class.Floating a) => Gen.Square (ShapeInt::+ShapeInt) a (Diagonal a)
genIdentityDiagonal =
   uncurry Block.Diagonal <$> Gen.stackDiagonal genIdentity genIdentity

genInvertibleDiagonal ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Gen.Square (ShapeInt::+ShapeInt) a (Diagonal a)
genInvertibleDiagonal =
   uncurry Block.Diagonal <$> Gen.stackDiagonal Gen.invertible Gen.invertible


type Square =
      Block.Square
         TypeFull () () TypeFull () () TypeFull () () TypeFull () ()
         Extent.Size Extent.Big Extent.Big ShapeInt ShapeInt

squareTransposeUnpack :: (Class.Floating a) => Square a -> Bool
squareTransposeUnpack a =
   Util.equalArray
      (Matrix.toFull (Matrix.transpose a))
      (Matrix.transpose (Matrix.toFull a))


_genSquareCompose, genSquare ::
   (Class.Floating a) => Gen.Square (ShapeInt::+ShapeInt) a (Square a)
_genSquareCompose =
   (\(a,b) (c,d) -> Block.Square a b c d) <$>
   ((,) <$> Gen.square <|||> Gen.matrix)
   <===>
   ((,) <$> Gen.matrix <|||> Gen.square)
genSquare =
   (\m -> case Square.split m of (a,b,c,d) -> Block.Square a b c d)
   <$>
   Gen.square

genIdentitySquare ::
   (Class.Floating a) => Gen.Square (ShapeInt::+ShapeInt) a (Square a)
genIdentitySquare =
   (\m -> case Square.split m of (a,b,c,d) -> Block.Square a b c d)
   <$>
   genIdentity

genInvertibleSquare ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Gen.Square (ShapeInt::+ShapeInt) a (Square a)
genInvertibleSquare =
   Gen.condition
      (\m@(Block.Square _a _b _c d) ->
         Util.invertible d && Util.invertible (Block.schurComplement m))
      genSquare


type Above =
      Block.Above TypeFull () () TypeFull () () Big ShapeInt ShapeInt ShapeInt

aboveTransposeUnpack :: (Class.Floating a) => Above a -> Bool
aboveTransposeUnpack a =
   Util.equalArray
      (Matrix.toFull (Matrix.transpose a))
      (Matrix.transpose (Matrix.toFull a))

genAbove ::
   (Class.Floating a) => Gen.Matrix (ShapeInt::+ShapeInt) ShapeInt a (Above a)
genAbove = Block.Above <$> Gen.matrix <===> Gen.matrix


type Beside =
      Block.Beside TypeFull () () TypeFull () () Big ShapeInt ShapeInt ShapeInt

besideTransposeUnpack :: (Class.Floating a) => Beside a -> Bool
besideTransposeUnpack a =
   Util.equalArray
      (Matrix.toFull (Matrix.transpose a))
      (Matrix.transpose (Matrix.toFull a))

genBeside ::
   (Class.Floating a) => Gen.Matrix ShapeInt (ShapeInt::+ShapeInt) a (Beside a)
genBeside = Block.Beside <$> Gen.matrix <|||> Gen.matrix


type Upper =
      Block.UpperTriangular TypeFull () () TypeFull () () TypeFull () ()
         MatrixShape.Filled ShapeInt ShapeInt

upperTransposeUnpack :: (Class.Floating a) => Upper a -> Bool
upperTransposeUnpack a =
   Util.equalArray
      (Matrix.toFull (Matrix.transpose a))
      (Matrix.transpose (Matrix.toFull a))


genUpper ::
   (Class.Floating a) => Gen.Square (ShapeInt::+ShapeInt) a (Upper a)
genUpper =
   (\m -> case Square.split m of (a,b,_c,d) -> Block.Upper a b d)
   <$>
   Gen.square

genIdentityUpper ::
   (Class.Floating a) => Gen.Square (ShapeInt::+ShapeInt) a (Upper a)
genIdentityUpper =
   (\m -> case Square.split m of (a,b,_c,d) -> Block.Upper a b d)
   <$>
   genIdentity

genInvertibleUpper ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Gen.Square (ShapeInt::+ShapeInt) a (Upper a)
genInvertibleUpper =
   Gen.condition
      (\(Block.Upper a _ d) -> Util.invertible a && Util.invertible d)
      genUpper


type Lower =
      Block.LowerTriangular TypeFull () () TypeFull () () TypeFull () ()
         MatrixShape.Filled ShapeInt ShapeInt

lowerTransposeUnpack :: (Class.Floating a) => Lower a -> Bool
lowerTransposeUnpack a =
   Util.equalArray
      (Matrix.toFull (Matrix.transpose a))
      (Matrix.transpose (Matrix.toFull a))


genLower ::
   (Class.Floating a) => Gen.Square (ShapeInt::+ShapeInt) a (Lower a)
genLower =
   (\m -> case Square.split m of (a,_b,c,d) -> Block.Lower a c d)
   <$>
   Gen.square

genIdentityLower ::
   (Class.Floating a) => Gen.Square (ShapeInt::+ShapeInt) a (Lower a)
genIdentityLower =
   (\m -> case Square.split m of (a,_b,c,d) -> Block.Lower a c d)
   <$>
   genIdentity

genInvertibleLower ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Gen.Square (ShapeInt::+ShapeInt) a (Lower a)
genInvertibleLower =
   Gen.condition
      (\(Block.Lower a _ d) -> Util.invertible a && Util.invertible d)
      genLower


type Symmetric =
      Block.Symmetric
         TypeFull () () TypeFull () () TypeFull () ()
         ShapeInt ShapeInt

symmetricTransposeUnpack :: (Class.Floating a) => Symmetric a -> Bool
symmetricTransposeUnpack a =
   Util.equalArray
      (Matrix.toFull (Matrix.transpose a))
      (Matrix.transpose (Matrix.toFull a))


genSymmetric ::
   (Class.Floating a) => Gen.Square (ShapeInt::+ShapeInt) a (Symmetric a)
genSymmetric =
   (\m -> case Square.split m of (a,b,_c,d) -> Block.Symmetric a b d)
   <$>
   Gen.square


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


testsVar ::
   (Show a, Class.Floating a, Eq a, RealOf a ~ ar, Class.Real ar, Show ar) =>
   [(String, Tagged a QC.Property)]
testsVar =
   prefix "Diagonal"
      (("transposeUnpack", checkForAll genDiagonal diagonalTransposeUnpack) :
       Generic.testsDistributive genDiagonal ++
       Multiply.testsVar genIdentityDiagonal genDiagonal ++
       Divide.testsVar genInvertibleDiagonal) ++

   prefix "Square"
      (("transposeUnpack", checkForAll genSquare squareTransposeUnpack) :
       Generic.testsDistributive genSquare ++
       Multiply.testsVar genIdentitySquare genSquare ++
       Divide.testsVar genInvertibleSquare) ++

   prefix "Above"
      (("transposeUnpack", checkForAll genAbove aboveTransposeUnpack) :
       Generic.testsDistributive genAbove ++
       Multiply.testsGeneralVar genAbove) ++

   prefix "Beside"
      (("transposeUnpack", checkForAll genBeside besideTransposeUnpack) :
       Generic.testsDistributive genBeside ++
       Multiply.testsGeneralVar genBeside) ++

   prefix "Upper"
      (("transposeUnpack", checkForAll genUpper upperTransposeUnpack) :
       Generic.testsDistributive genUpper ++
       Multiply.testsVar genIdentityUpper genUpper ++
       Divide.testsVar genInvertibleUpper) ++

   prefix "Lower"
      (("transposeUnpack", checkForAll genLower lowerTransposeUnpack) :
       Generic.testsDistributive genLower ++
       Multiply.testsVar genIdentityLower genLower ++
       Divide.testsVar genInvertibleLower) ++

   prefix "Symmetric"
      (("transposeUnpack", checkForAll genSymmetric symmetricTransposeUnpack) :
       Generic.testsDistributive genSymmetric) ++
   []