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

import qualified Test.Divide as Divide
import qualified Test.Multiply as Multiply
import qualified Test.Generic as Generic
import qualified Test.Indexed as Indexed
import qualified Test.Generator as Gen
import qualified Test.Utility as Util
import Test.Generator ((<-*#>), (<#*|>), (<.*#>), (<#*#>), (<#\#>), (<#=#>))
import Test.Utility
         (approxReal, approxArray, approxArrayTol, approxMatrix,
          approxVector, equalArray, Tagged, genOrder, (!===))

import qualified Numeric.LAPACK.Orthogonal.Householder as HH
import qualified Numeric.LAPACK.Matrix.HermitianPositiveDefinite as HermitianPD
import qualified Numeric.LAPACK.Matrix.Hermitian as Hermitian
import qualified Numeric.LAPACK.Matrix.Triangular as Triangular
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 as Matrix
import qualified Numeric.LAPACK.Vector as Vector
import Numeric.LAPACK.Matrix.Hermitian (Hermitian)
import Numeric.LAPACK.Matrix.Square (Square)
import Numeric.LAPACK.Matrix.Shape (Order)
import Numeric.LAPACK.Matrix
         (General, ShapeInt, (#+#), (-*#), (##*#), (#*##), (#*|), (|||))
import Numeric.LAPACK.Vector (Vector, (.*|))
import Numeric.LAPACK.Scalar (RealOf, selectReal)

import qualified Numeric.Netlib.Class as Class

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

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

import qualified Data.NonEmpty.Class as NonEmptyC
import qualified Data.NonEmpty as NonEmpty
import Data.Semigroup ((<>))
import Data.Tuple.HT (uncurry3, mapFst)

import qualified Test.QuickCheck as QC



generalFromHermitian ::
   (Shape.C sh, Class.Floating a) => Hermitian sh a -> General sh sh a
generalFromHermitian = Matrix.fromFull . Hermitian.toSquare

stack ::
   (Class.Floating a) =>
   (Hermitian ShapeInt a, General ShapeInt ShapeInt a, Hermitian ShapeInt a) ->
   Bool
stack (a,b,c) =
   let abc = generalFromHermitian $ Hermitian.stack a b c
   in equalArray abc $
         (Matrix.fromFull (Hermitian.toSquare a) ||| b
          !===
          Matrix.adjoint b ||| Matrix.fromFull (Hermitian.toSquare c))

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

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

gramianAdjoint ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   General ShapeInt ShapeInt a -> Bool
gramianAdjoint x =
   approxArray
      (generalFromHermitian $ Hermitian.gramianAdjoint x)
      (Matrix.adaptOrder x $ x <> Matrix.adjoint x)

gramianNonAdjoint ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   General ShapeInt ShapeInt a -> Bool
gramianNonAdjoint x =
   approxArray
      (Matrix.forceOrder (ArrMatrix.shapeOrder $ ArrMatrix.shape x) $
       Hermitian.gramian $ Matrix.adjoint x)
      (Hermitian.gramianAdjoint x)

congruenceDiagonal ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Vector ShapeInt ar, General ShapeInt ShapeInt a) -> Bool
congruenceDiagonal (d,a) =
   approxArray
      (generalFromHermitian $ Hermitian.congruenceDiagonal d a)
      (Matrix.adjoint a <> Matrix.scaleRowsReal d a)

congruenceDiagonalAdjoint ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (General ShapeInt ShapeInt a, Vector ShapeInt ar) -> Bool
congruenceDiagonalAdjoint (a,d) =
   approxMatrix 1e-5
      (generalFromHermitian $ Hermitian.congruenceDiagonalAdjoint a d)
      (Matrix.scaleColumnsReal d a <> Matrix.adjoint a)

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

congruence ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Hermitian ShapeInt a, General ShapeInt ShapeInt a) -> Bool
congruence (b,a) =
   approxArray
      (Hermitian.toSquare $ Hermitian.congruence b a)
      (Square.congruence (Hermitian.toSquare b) a)

congruenceAdjoint ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (General ShapeInt ShapeInt a, Hermitian ShapeInt a) -> Bool
congruenceAdjoint (a,b) =
   approxMatrix 1e-5
      (Hermitian.toSquare $ Hermitian.congruenceAdjoint a b)
      (Square.congruenceAdjoint a $ Hermitian.toSquare b)

congruenceCongruenceDiagonal ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Order -> (Vector ShapeInt ar, General ShapeInt ShapeInt a) -> Bool
congruenceCongruenceDiagonal order (d,a) =
   approxArray
      (Hermitian.congruenceDiagonal d a)
      (Hermitian.congruence (Hermitian.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
      (generalFromHermitian $ Hermitian.anticommutator a b)
      ((Matrix.adjoint b <> a) #+# (Matrix.adjoint 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
      (Hermitian.anticommutator a b)
      (Hermitian.anticommutator b a)

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


outer ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Order -> Vector ShapeInt a -> Bool
outer order x =
   approxArray
      (generalFromHermitian $ Hermitian.outer order x)
      (Matrix.outer order x x)


genScaledVectors ::
   (NonEmptyC.Gen f, Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Gen.VectorInt a (ShapeInt, f (ar, Vector ShapeInt a))
genScaledVectors = Gen.listOfVector ((,) <$> Gen.scalarReal <.*#> Gen.vector)

sumRank1 ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Order -> (ShapeInt, [(ar, Vector ShapeInt a)]) -> Bool
sumRank1 order (sh,xs) =
   approxArray
      (generalFromHermitian $ Hermitian.sumRank1 order sh xs)
      (Util.addMatrices (MatrixShape.general order sh sh) $
       fmap (rank1 order) xs)

sumRank1NonEmpty ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Order -> NonEmpty.T [] (ar, Vector ShapeInt a) -> Bool
sumRank1NonEmpty order xs =
   approxArray
      (generalFromHermitian $ Hermitian.sumRank1NonEmpty order xs)
      (NonEmpty.foldl1 (ArrMatrix.lift2 Vector.add) $ fmap (rank1 order) xs)

rank1 ::
   (Eq size, Shape.C size, Class.Floating a) =>
   Order -> (RealOf a, Vector size a) -> Matrix.General size size a
rank1 order (r,x) = Matrix.scaleReal r $ Matrix.outer order x x


genScaledVectorPairs ::
   (NonEmptyC.Gen f, Class.Floating a) =>
   Gen.VectorInt a (ShapeInt, f (a, (Vector ShapeInt a, Vector ShapeInt a)))
genScaledVectorPairs =
   flip Gen.mapGen Gen.vectorDim $ \maxElem size ->
      fmap ((,) size) $
      NonEmptyC.genOf $
         liftA2 (,) (Util.genElement maxElem) $
         liftA2 (,) (Util.genVector maxElem size) (Util.genVector maxElem size)

sumRank2 ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Order -> (ShapeInt, [(a, (Vector ShapeInt a, Vector ShapeInt a))]) -> Bool
sumRank2 order (sh,xys) =
   approxArray
      (generalFromHermitian $ Hermitian.sumRank2 order sh xys)
      (Util.addMatrices (MatrixShape.general order sh sh) $
       fmap (rank2 order) xys)

sumRank2NonEmpty ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Order -> NonEmpty.T [] (a, (Vector ShapeInt a, Vector ShapeInt a)) -> Bool
sumRank2NonEmpty order xys =
   approxArray
      (generalFromHermitian $ Hermitian.sumRank2NonEmpty order xys)
      (NonEmpty.foldl1 (ArrMatrix.lift2 Vector.add) $ fmap (rank2 order) xys)

rank2 ::
   (Eq size, Shape.C size, Class.Floating a) =>
   Order -> (a, (Vector size a, Vector size a)) -> Matrix.General size size a
rank2 order (a,(x,y)) =
   let ax = a.*|x
   in ArrMatrix.lift2 Vector.add
         (Matrix.outer order ax y)
         (Matrix.outer order y ax)


addAdjoint ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Square ShapeInt a -> Bool
addAdjoint x =
   approxArray
      (Hermitian.toSquare $ Hermitian.addAdjoint x)
      (Matrix.adjoint x #+# x)



multiplyVectorLeft ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Vector ShapeInt a, Hermitian ShapeInt a) -> Bool
multiplyVectorLeft (x,a) =
   approxVector (x -*# Hermitian.toSquare a) (x -*# a)

multiplyVectorRight ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Hermitian ShapeInt a, Vector ShapeInt a) -> Bool
multiplyVectorRight (a,x) =
   approxVector (Hermitian.toSquare a #*| x) (a #*| x)


multiplyLeft ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (General ShapeInt ShapeInt a, Hermitian ShapeInt a) -> Bool
multiplyLeft (a,b) =
   approxMatrix 1e-5 (a ##*# Hermitian.toSquare b) (a ##*# b)

multiplyRight ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Hermitian ShapeInt a, General ShapeInt ShapeInt a) -> Bool
multiplyRight (a,b) =
   approxArray (Hermitian.toSquare a #*## b) (a #*## b)


choleskyQR ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Matrix.Tall ShapeInt ShapeInt a -> QC.Property
choleskyQR a =
   let qr = HH.fromMatrix a
       r = HH.tallExtractR qr
   in HH.determinantAbsolute qr > 0.1
      QC.==>
      approxArrayTol 1e-1
         (Matrix.scaleRows (Array.map signum $ Triangular.takeDiagonal r) $
          Triangular.toSquare r)
         (Triangular.toSquare $
          HermitianPD.decompose $ Hermitian.gramian $ Matrix.fromFull a)


genInvertible ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Gen.MatrixInt a (Hermitian ShapeInt a)
genInvertible = Gen.condition Util.invertible Gen.hermitian

inverse ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Hermitian ShapeInt a -> Bool
inverse a =
   approxArrayTol
      (selectReal 1 1e-5)
      (Hermitian.toSquare $ Hermitian.inverse a)
      (Square.inverse $ Hermitian.toSquare a)


solve ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Hermitian ShapeInt a, Matrix.General ShapeInt ShapeInt a) -> Bool
solve (a, b) =
   approxMatrix (selectReal 1 1e-5)
      (Hermitian.solve a b)
      (Square.solve (Hermitian.toSquare a) b)



genPositiveDefinite ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Gen.MatrixInt a (Hermitian ShapeInt a)
genPositiveDefinite =
   Hermitian.gramian . Matrix.fromFull <$> Gen.gramian Gen.fullRankTall

determinantPD ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Hermitian ShapeInt a -> Bool
determinantPD a =
   approxReal (selectReal 100 1e-4)
      (Hermitian.determinant a)
      (HermitianPD.determinant a)

inversePD ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Hermitian ShapeInt a -> Bool
inversePD a =
   approxArrayTol (selectReal 1000 1e-4)
      (Hermitian.inverse a)
      (HermitianPD.inverse a)

solvePD ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Hermitian ShapeInt a, Matrix.General ShapeInt ShapeInt a) -> Bool
solvePD (a,b) =
   approxArrayTol (selectReal 1000 1e-4)
      (Hermitian.solve a b)
      (HermitianPD.solve a b)

solveDecomposedPD ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Hermitian ShapeInt a, Matrix.General ShapeInt ShapeInt a) -> Bool
solveDecomposedPD (a,b) =
   approxArrayTol (selectReal 1e-1 1e-6)
      (HermitianPD.solve a b)
      (HermitianPD.solveDecomposed (HermitianPD.decompose a) b)



eigenvaluesDeterminant ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Hermitian ShapeInt a -> Bool
eigenvaluesDeterminant a =
   approxReal
      (selectReal 1e-1 1e-5)
      (Hermitian.determinant a)
      (Vector.product $ Hermitian.eigenvalues a)

eigensystem ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Hermitian ShapeInt a -> Bool
eigensystem a =
   approxMatrix 1e-4 a $
   uncurry Hermitian.congruenceDiagonalAdjoint $
   mapFst Matrix.fromFull $ Hermitian.eigensystem a



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, Show ar, Class.Floating a, Eq a, RealOf a ~ ar, Class.Real ar) =>
   [(String, Tagged a QC.Property)]
testsVar =
   ("index",
      checkForAll (Indexed.genMatrixIndex Gen.hermitian) Indexed.unitDot) :
   ("forceOrder",
      checkForAllExtra genOrder
         ((,) <$> Gen.hermitian <#*|> Gen.vector) Generic.forceOrder) :
   ("addDistributive",
      checkForAll
         (Generic.genDistribution Gen.hermitian)
         Generic.addDistributive) :
   ("subDistributive",
      checkForAll
         (Generic.genDistribution Gen.hermitian)
         Generic.subDistributive) :

   ("stack",
      checkForAll (Gen.stack3 Gen.hermitian Gen.matrix Gen.hermitian) stack) :
   ("split",
      checkForAll Gen.hermitian split) :
   ("gramian",
      checkForAll Gen.matrix gramian) :
   ("gramianAdjoint",
      checkForAll Gen.matrix gramianAdjoint) :
   ("gramianNonAdjoint",
      checkForAll Gen.matrix gramianNonAdjoint) :
   ("congruenceDiagonal",
      checkForAll ((,) <$> Gen.vectorReal <-*#> Gen.matrix) congruenceDiagonal) :
   ("congruence",
      checkForAll ((,) <$> Gen.hermitian <#*#> Gen.matrix) congruence) :
   ("congruenceDiagonalAdjoint",
      checkForAll
         ((,) <$> Gen.matrix <#*|> Gen.vectorReal) congruenceDiagonalAdjoint) :
   ("congruenceDiagonalGramian",
      checkForAll Gen.matrix congruenceDiagonalGramian) :
   ("congruenceAdjoint",
      checkForAll ((,) <$> Gen.matrix <#*#> Gen.hermitian) congruenceAdjoint) :
   ("congruenceCongruenceDiagonal",
      checkForAllExtra genOrder
         ((,) <$>
            Gen.vectorReal <-*#> Gen.matrix) congruenceCongruenceDiagonal) :
   ("anticommutator",
      checkForAll ((,) <$> Gen.matrix <#=#> Gen.matrix) anticommutator) :
   ("anticommutatorCommutative",
      checkForAll ((,) <$> Gen.matrix <#=#> Gen.matrix)
         anticommutatorCommutative) :
   ("anticommutatorAdjoint",
      checkForAll ((,) <$> Gen.matrix <#=#> Gen.matrix) anticommutatorAdjoint) :
   ("outer",
      checkForAllExtra genOrder Gen.vector outer) :
   ("sumRank1",
      checkForAllExtra genOrder genScaledVectors sumRank1) :
   ("sumRank1NonEmpty",
      checkForAllExtra genOrder (snd <$> genScaledVectors) sumRank1NonEmpty) :
   ("sumRank2",
      checkForAllExtra genOrder genScaledVectorPairs sumRank2) :
   ("sumRank2NonEmpty",
      checkForAllExtra genOrder
         (snd <$> genScaledVectorPairs) sumRank2NonEmpty) :
   ("addAdjoint",
      checkForAll Gen.square addAdjoint) :
   ("multiplySquare",
      checkForAll Gen.hermitian Multiply.multiplySquare) :
   ("squareSquare",
      checkForAll Gen.hermitian Multiply.squareSquare) :
   ("power",
      checkForAllExtra (QC.choose (0,10)) Gen.hermitian Multiply.power) :

   ("multiplyVectorLeft",
      checkForAll ((,) <$> Gen.vector <-*#> Gen.hermitian) multiplyVectorLeft) :
   ("multiplyVectorRight",
      checkForAll ((,) <$> Gen.hermitian <#*|> Gen.vector) multiplyVectorRight) :
   ("multiplyLeft",
      checkForAll ((,) <$> Gen.matrix <#*#> Gen.hermitian) multiplyLeft) :
   ("multiplyRight",
      checkForAll ((,) <$> Gen.hermitian <#*#> Gen.matrix) multiplyRight) :

   ("determinant",
      checkForAll Gen.hermitian Divide.determinant) :
   ("choleskyQR",
      checkForAll Gen.tall choleskyQR) :

   ("inverse",
      checkForAll genInvertible inverse) :
   ("solve",
      checkForAll ((,) <$> genInvertible <#\#> Gen.matrix) solve) :
   Divide.testsVar genInvertible ++

   ("determinantPD",
      checkForAll genPositiveDefinite determinantPD) :
   ("inversePD",
      checkForAll genPositiveDefinite inversePD) :
   ("solvePD",
      checkForAll ((,) <$> genPositiveDefinite <#\#> Gen.matrix) solvePD) :
   ("solveDecomposedPD",
      checkForAll
         ((,) <$> genPositiveDefinite <#\#> Gen.matrix) solveDecomposedPD) :

   ("eigenvaluesDeterminant",
      checkForAll Gen.hermitian eigenvaluesDeterminant) :
   ("eigensystem",
      checkForAll Gen.hermitian eigensystem) :
   []