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

import qualified Test.Generator as Gen
import qualified Test.Utility as Util
import Test.Generator ((<-*#>), (<#*|>), (<#*#>), (<#=#>))
import Test.Utility
         (approxArray, approxMatrix, equalArray, Tagged, genOrder, (!===))

import qualified Numeric.LAPACK.Matrix.Symmetric as Symmetric
import qualified Numeric.LAPACK.Matrix.Triangular as Triangular
import qualified Numeric.LAPACK.Matrix.Array as ArrMatrix
import qualified Numeric.LAPACK.Matrix as Matrix
import qualified Numeric.LAPACK.Vector as Vector
import Numeric.LAPACK.Matrix.Symmetric (Symmetric)
import Numeric.LAPACK.Matrix.Shape (Order)
import Numeric.LAPACK.Matrix (General, ShapeInt, (#+#), (|||))
import Numeric.LAPACK.Vector (Vector)
import Numeric.LAPACK.Scalar (RealOf)

import qualified Numeric.Netlib.Class as Class

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

import Control.Applicative ((<$>))

import Data.Semigroup ((<>))
import Data.Tuple.HT (uncurry3)

import qualified Test.QuickCheck as QC



generalFromSymmetric ::
   (Shape.C sh, Class.Floating a) => Symmetric sh a -> General sh sh a
generalFromSymmetric = Matrix.fromFull . Triangular.toSquare

stack ::
   (Class.Floating a) =>
   (Symmetric ShapeInt a, General ShapeInt ShapeInt a, Symmetric ShapeInt a) ->
   Bool
stack (a,b,c) =
   let abc = generalFromSymmetric $ Symmetric.stack a b c
   in equalArray abc $
         (Matrix.fromFull (Triangular.toSquare a) ||| b
          !===
          Matrix.transpose b ||| Matrix.fromFull (Triangular.toSquare c))

split :: (Class.Floating a) => Symmetric (ShapeInt:+:ShapeInt) a -> Bool
split abc = equalArray abc $ uncurry3 Symmetric.stack $ Symmetric.split abc


gramian ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   General ShapeInt ShapeInt a -> Bool
gramian x =
   approxArray
      (generalFromSymmetric $ Symmetric.gramian x)
      (Matrix.transpose x <> x)

gramianTransposed ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   General ShapeInt ShapeInt a -> Bool
gramianTransposed x =
   approxArray
      (generalFromSymmetric $ Symmetric.gramianTransposed x)
      (Matrix.adaptOrder x $ x <> Matrix.transpose x)

gramianNonTransposed ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   General ShapeInt ShapeInt a -> Bool
gramianNonTransposed x =
   approxArray
      (Matrix.forceOrder (ArrMatrix.shapeOrder $ ArrMatrix.shape x) $
       Symmetric.gramian $ Matrix.transpose x)
      (Symmetric.gramianTransposed x)

congruenceDiagonal ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Vector ShapeInt a, General ShapeInt ShapeInt a) -> Bool
congruenceDiagonal (d,a) =
   approxArray
      (generalFromSymmetric $ Symmetric.congruenceDiagonal d a)
      (Matrix.transpose a <> Matrix.scaleRows d a)

congruenceDiagonalTransposed ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (General ShapeInt ShapeInt a, Vector ShapeInt a) -> Bool
congruenceDiagonalTransposed (a,d) =
   approxMatrix 1e-5
      (generalFromSymmetric $ Symmetric.congruenceDiagonalTransposed a d)
      (Matrix.scaleColumns d a <> Matrix.transpose a)

congruenceDiagonalGramian ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   General ShapeInt ShapeInt a -> Bool
congruenceDiagonalGramian a =
   approxArray
      (Symmetric.congruenceDiagonal (Vector.one $ Matrix.height a) a)
      (Symmetric.gramian a)

congruence ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Symmetric ShapeInt a, General ShapeInt ShapeInt a) -> Bool
congruence (b,a) =
   approxArray
      (generalFromSymmetric $ Symmetric.congruence b a)
      (Matrix.transpose a <> generalFromSymmetric b <> a)

congruenceTransposed ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (General ShapeInt ShapeInt a, Symmetric ShapeInt a) -> Bool
congruenceTransposed (a,b) =
   approxMatrix 1e-5
      (generalFromSymmetric $ Symmetric.congruenceTransposed a b)
      (a <> generalFromSymmetric b <> Matrix.transpose a)

congruenceCongruenceDiagonal ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Order -> (Vector ShapeInt a, General ShapeInt ShapeInt a) -> Bool
congruenceCongruenceDiagonal order (d,a) =
   approxArray
      (Symmetric.congruenceDiagonal d a)
      (Symmetric.congruence (Triangular.diagonal order d) a)

anticommutator ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (General ShapeInt ShapeInt a, General ShapeInt ShapeInt a) -> Bool
anticommutator (a,b) =
   approxArray
      (generalFromSymmetric $ Symmetric.anticommutator a b)
      ((Matrix.transpose b <> a) #+# (Matrix.transpose a <> b))

anticommutatorCommutative ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (General ShapeInt ShapeInt a, General ShapeInt ShapeInt a) -> Bool
anticommutatorCommutative (a,b) =
   approxMatrix 1e-5
      (Symmetric.anticommutator a b)
      (Symmetric.anticommutator b a)

anticommutatorTransposed ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (General ShapeInt ShapeInt a, General ShapeInt ShapeInt a) -> Bool
anticommutatorTransposed (a,b) =
   approxArray
      (Matrix.forceOrder (ArrMatrix.shapeOrder $ ArrMatrix.shape b) $
       Symmetric.anticommutator (Matrix.transpose a) (Matrix.transpose b))
      (Symmetric.anticommutatorTransposed a b)



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)

checkForAllExtra ::
   (Show a, Show b, QC.Testable test) =>
   QC.Gen a -> Gen.T dim tag 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 =
   ("stack",
      checkForAll (Gen.stack3 Gen.triangular Gen.matrix Gen.triangular) stack) :
   ("split",
      checkForAll Gen.triangular split) :

   ("gramian",
      checkForAll Gen.matrix gramian) :
   ("gramianTransposed",
      checkForAll Gen.matrix gramianTransposed) :
   ("gramianNonTransposed",
      checkForAll Gen.matrix gramianNonTransposed) :
   ("congruenceDiagonal",
      checkForAll ((,) <$> Gen.vector <-*#> Gen.matrix) congruenceDiagonal) :
   ("congruence",
      checkForAll ((,) <$> Gen.triangular <#*#> Gen.matrix) congruence) :
   ("congruenceDiagonalTransposed",
      checkForAll ((,) <$> Gen.matrix <#*|> Gen.vector)
         congruenceDiagonalTransposed) :
   ("congruenceDiagonalGramian",
      checkForAll Gen.matrix congruenceDiagonalGramian) :
   ("congruenceTransposed",
      checkForAll ((,) <$> Gen.matrix <#*#> Gen.triangular)
         congruenceTransposed) :
   ("congruenceCongruenceDiagonal",
      checkForAllExtra genOrder
         ((,) <$> Gen.vector <-*#> Gen.matrix) congruenceCongruenceDiagonal) :
   ("anticommutator",
      checkForAll ((,) <$> Gen.matrix <#=#> Gen.matrix) anticommutator) :
   ("anticommutatorCommutative",
      checkForAll ((,) <$> Gen.matrix <#=#> Gen.matrix)
         anticommutatorCommutative) :
   ("anticommutatorTransposed",
      checkForAll ((,) <$> Gen.matrix <#=#> Gen.matrix)
         anticommutatorTransposed) :
   []