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

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

import qualified Numeric.LAPACK.Matrix.Symmetric as Symmetric
import qualified Numeric.LAPACK.Matrix.Array as ArrMatrix
import qualified Numeric.LAPACK.Matrix.Shape.Omni as Omni
import qualified Numeric.LAPACK.Matrix.Shape as MatrixShape
import qualified Numeric.LAPACK.Matrix.Layout as Layout
import qualified Numeric.LAPACK.Matrix as Matrix
import qualified Numeric.LAPACK.Vector as Vector
import Numeric.LAPACK.Matrix.Layout (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 qualified Data.NonEmpty.Class as NonEmptyC
import qualified Data.NonEmpty as NonEmpty
import qualified Data.List as List
import Data.Semigroup ((<>))
import Data.Tuple.HT (uncurry3)

import qualified Test.QuickCheck as QC



type SymmetricP pack sh = ArrMatrix.FullQuadratic pack Omni.Symmetric sh

genSymmetric ::
   (Logic.Dim sh, Shape.Indexed sh, Shape.Index sh ~ ix, Eq ix,
    Class.Floating a) =>
   Layout.PackingSingleton pack ->
   Gen.Square sh a (SymmetricP pack sh a)
genSymmetric p = repack p <$> Gen.symmetric

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

stack ::
   (Layout.Packing pack, Class.Floating a) =>
   (SymmetricP pack ShapeInt a,
    General ShapeInt ShapeInt a,
    SymmetricP pack ShapeInt a) ->
   Bool
stack (a,b,c) =
   equalArray
      (generalFromSymmetric $ Symmetric.stack a b c)
      (generalFromSymmetric a ||| b
       !===
       Matrix.transpose b ||| generalFromSymmetric c)

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


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

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

gramianNonTransposed ::
   (Layout.Packing pack) =>
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Layout.PackingSingleton pack ->
   General ShapeInt ShapeInt a -> Bool
gramianNonTransposed pack x =
   approxArray
      (Matrix.forceOrder (ArrMatrix.order x) $
       Symmetric.gramian $ Matrix.transpose x)
      (ArrMatrix.requirePacking pack $ Symmetric.gramianTransposed x)

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

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

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

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

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

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

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

anticommutatorCommutative ::
   (Layout.Packing pack) =>
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Layout.PackingSingleton pack ->
   (General ShapeInt ShapeInt a, General ShapeInt ShapeInt a) -> Bool
anticommutatorCommutative pack (a,b) =
   approxMatrix 1e-5
      (ArrMatrix.requirePacking pack $
       Symmetric.anticommutator a b)
      (Symmetric.anticommutator b a)

anticommutatorTransposed ::
   (Layout.Packing pack) =>
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Layout.PackingSingleton pack ->
   (General ShapeInt ShapeInt a, General ShapeInt ShapeInt a) -> Bool
anticommutatorTransposed pack (a,b) =
   approxArray
      (Matrix.forceOrder (ArrMatrix.order b) $
       Symmetric.anticommutator (Matrix.transpose a) (Matrix.transpose b))
      (ArrMatrix.requirePacking pack $ Symmetric.anticommutatorTransposed a b)


tensorProduct ::
   (Layout.Packing pack) =>
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Layout.PackingSingleton pack ->
   Order -> Vector ShapeInt a -> Bool
tensorProduct pack order x =
   approxArray
      (generalFromSymmetric $
       ArrMatrix.requirePacking pack $ Symmetric.tensorProduct order x)
      (Matrix.tensorProduct order x x)


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

sumRank1 ::
   (Layout.Packing pack) =>
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Layout.PackingSingleton pack ->
   Order -> (ShapeInt, [(a, Vector ShapeInt a)]) -> Bool
sumRank1 pack order (sh,xs) =
   approxArray
      (generalFromSymmetric $
       ArrMatrix.requirePacking pack $ Symmetric.sumRank1 order sh xs)
      (Util.addMatrices (MatrixShape.general order sh sh) $
       fmap (rank1 order) xs)

sumRank1NonEmpty ::
   (Layout.Packing pack) =>
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Layout.PackingSingleton pack ->
   Order -> NonEmpty.T [] (a, Vector ShapeInt a) -> Bool
sumRank1NonEmpty pack order xs =
   approxArray
      (generalFromSymmetric $
       ArrMatrix.requirePacking pack $ Symmetric.sumRank1NonEmpty order xs)
      (NonEmpty.foldl1 (ArrMatrix.lift2 Vector.add) $ fmap (rank1 order) xs)

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


addTransposed ::
   (Layout.Packing pack) =>
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Layout.PackingSingleton pack ->
   Matrix.Square ShapeInt a -> Bool
addTransposed pack x =
   approxArray
      (Symmetric.toSquare $ ArrMatrix.requirePacking pack $
         Symmetric.addTransposed x)
      (Matrix.transpose x #+# x)



genInvertible ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Layout.PackingSingleton pack ->
   Gen.MatrixInt a (SymmetricP pack ShapeInt a)
genInvertible pack =
   repack pack <$> Gen.condition Util.invertible Gen.symmetric



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 =
   concat $
   List.transpose
      [Util.suffix "Packed"   (testsVarPacking Layout.Packed),
       Util.suffix "Unpacked" (testsVarPacking Layout.Unpacked)]

testsVarPacking ::
   (Layout.Packing pack) =>
   (Show a, Class.Floating a, Eq a, RealOf a ~ ar, Class.Real ar) =>
   Layout.PackingSingleton pack -> [(String, Tagged a QC.Property)]
testsVarPacking p =
   ("index",
      checkForAll (Indexed.genMatrixIndex $ genSymmetric p) Indexed.unitDot) :
   ("forceOrder",
      checkForAllExtra genOrder
         ((,) <$> genSymmetric p <#*|> Gen.vector) Generic.forceOrder) :
   ("forceOrderInverse",
      checkForAll (genSymmetric p) Generic.forceOrderInverse) :
   Generic.testsDistributive (Gen.asMatrixInt $ genSymmetric p) ++

   ("stack",
      checkForAll
         (Gen.stack3 (genSymmetric p) Gen.matrix (genSymmetric p))
         stack) :
   ("split",
      checkForAll (genSymmetric p) split) :

   ("gramian",
      checkForAll Gen.matrix (gramian p)) :
   ("gramianTransposed",
      checkForAll Gen.matrix (gramianTransposed p)) :
   ("gramianNonTransposed",
      checkForAll Gen.matrix (gramianNonTransposed p)) :
   ("congruenceDiagonal",
      checkForAll ((,) <$> Gen.vector <-*#> Gen.matrix)
         (congruenceDiagonal p)) :
   ("congruence",
      checkForAll ((,) <$> genSymmetric p <#*#> Gen.matrix) congruence) :
   ("congruenceDiagonalTransposed",
      checkForAll ((,) <$> Gen.matrix <#*|> Gen.vector)
         (congruenceDiagonalTransposed p)) :
   ("congruenceDiagonalGramian",
      checkForAll Gen.matrix (congruenceDiagonalGramian p)) :
   ("congruenceTransposed",
      checkForAll ((,) <$> Gen.matrix <#*#> genSymmetric p)
         congruenceTransposed) :
   ("congruenceCongruenceDiagonal",
      checkForAllExtra genOrder
         ((,) <$> Gen.vector <-*#> Gen.matrix)
         (congruenceCongruenceDiagonal p)) :
   ("anticommutator",
      checkForAll ((,) <$> Gen.matrix <#=#> Gen.matrix) (anticommutator p)) :
   ("anticommutatorCommutative",
      checkForAll ((,) <$> Gen.matrix <#=#> Gen.matrix)
         (anticommutatorCommutative p)) :
   ("anticommutatorTransposed",
      checkForAll ((,) <$> Gen.matrix <#=#> Gen.matrix)
         (anticommutatorTransposed p)) :
   ("tensorProduct",
      checkForAllExtra genOrder Gen.vector (tensorProduct p)) :
   ("sumRank1",
      checkForAllExtra genOrder genScaledVectors (sumRank1 p)) :
   ("sumRank1NonEmpty",
      checkForAllExtra genOrder
         (snd <$> genScaledVectors) (sumRank1NonEmpty p)) :
   ("addTransposed",
      checkForAll Gen.square (addTransposed p)) :

   Mosaic.testsVar Mosaic.Symmetric p ++

   ("determinant",
      checkForAll (genSymmetric p) Divide.determinant) :
   ("solve",
      checkForAll ((,) <$> genInvertible p <#\#> Gen.matrix) Divide.solve) :
   ("solveIdentity",
      checkForAll
         ((,) <$>
            (repack p <$> Gen.identity `asTypeOf` Gen.symmetric)
               <#\#> Gen.matrix)
         Divide.solveIdentity) :
   ("inverse",
      checkForAll (genInvertible p) Divide.inverse) :
   Divide.testsVar (genInvertible p) ++
   []