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

import qualified Test.Divide as Divide
import qualified Test.Multiply as Multiply
import qualified Test.Generator as Gen
import qualified Test.Utility as Util
import Test.Generator ((<#*#>), (<#\#>))
import Test.Utility (approx, approxArray, approxArrayTol, approxMatrix, Tagged)

import qualified Numeric.LAPACK.Matrix.Square as Square
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.Square (Square)
import Numeric.LAPACK.Matrix
         (ShapeInt, (##*#), (#*##), (##/#), (#\##), (#*\), (\*#))
import Numeric.LAPACK.Scalar (RealOf, absolute, selectReal)

import qualified Numeric.Netlib.Class as Class

import Control.Applicative ((<$>))

import Data.Function.HT (nest)
import Data.Semigroup ((<>))
import Data.Complex (Complex)

import qualified Test.QuickCheck as QC


congruence ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Square ShapeInt a, Matrix.General ShapeInt ShapeInt a) -> Bool
congruence (b,a) =
   approxArray
      (Square.toFull $ Square.congruence b a)
      (Matrix.adjoint a <> (b #*## a))

congruenceAdjoint ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Matrix.General ShapeInt ShapeInt a, Square ShapeInt a) -> Bool
congruenceAdjoint (a,b) =
   approxArray
      (Square.toFull $ Square.congruenceAdjoint a b)
      (a <> b #*## Matrix.adjoint a)

multiplySquare ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Square ShapeInt a -> Bool
multiplySquare a =
   approxArray (Square.square a) (Square.multiply a a)

multiplyPower ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Int -> Square ShapeInt a -> Bool
multiplyPower n a =
   let b = Square.power (fromIntegral n) a
       c = nest n (Square.multiply a) $ Square.identityFrom a
       normInf1 = Vector.normInf1 . ArrMatrix.toVector
   in approxArrayTol (1e-6 * (normInf1 b + normInf1 c)) b c


determinantSingleton ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   a -> Bool
determinantSingleton a =
   approx 1e-5 a (Square.determinant $ Square.autoFromList [a])

determinantTranspose ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Square ShapeInt a -> Bool
determinantTranspose a =
   approx 1e-5
      (Square.determinant a) (Square.determinant $ Square.transpose a)


multiplyDeterminant ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Square ShapeInt a, Square ShapeInt a) -> Bool
multiplyDeterminant (a,b) =
   let detA = Square.determinant a
       detB = Square.determinant b
   in approx
         (1e-2 * max 1 (absolute detA) * max 1 (absolute detB))
         (Square.determinant (a<>b))
         (detA * detB)

multiplyInverse ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Square ShapeInt a -> Bool
multiplyInverse a = Util.isIdentity 1e-4 $ Square.inverse a <> a


multiplySolve ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Square ShapeInt a, Matrix.General ShapeInt ShapeInt a) -> Bool
multiplySolve (a, b) =
   approxMatrix 1e-2 (a #*## a #\## b) b

schur ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Square ShapeInt a -> Bool
schur a =
   let (q,r) = Square.schur a
   in approxMatrix 1e-4 a $
      Square.congruenceAdjoint (Matrix.fromFull q) (Matrix.toFull r)

schurComplex :: (Class.Real a) => Square ShapeInt (Complex a) -> Bool
schurComplex a =
   let (q,r) = Square.schurComplex a
   in approxMatrix 1e-4 a $ q #*## r #*## Square.adjoint q


genDiagonalizable ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Gen.MatrixInt a (Square ShapeInt a)
genDiagonalizable = flip Gen.mapQC Gen.invertible $ \a -> do
   d <- Util.genDistinct [-3..3] [-10..10] (Square.size a)
   return $ a #\## d \*# a


eigensystem ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Square ShapeInt a -> Bool
eigensystem a =
   let (vr,d,vlAdj) = Square.eigensystem a
       scal = Square.takeDiagonal $ vlAdj <> vr
   in  approxMatrix (selectReal 1e-1 1e-5)
         (Matrix.toComplex a)
         (vr #*\ Vector.divide d scal ##*# vlAdj)

eigensystemLeft ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Square ShapeInt a -> Bool
eigensystemLeft a =
   let (_vr,d,vlAdj) = Square.eigensystem a
   in  approxMatrix (selectReal 1e-1 1e-5)
         (Matrix.toComplex a)
         (vlAdj #\## d \*# vlAdj)

eigensystemRight ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Square ShapeInt a -> Bool
eigensystemRight a =
   let (vr,d,_vlAdj) = Square.eigensystem a
   in  approxMatrix (selectReal 1e-1 1e-5)
         (Matrix.toComplex a)
         (vr #*\ d ##/# vr)



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, Show ar, Class.Real ar) =>
   [(String, Tagged a QC.Property)]
testsVar =
   ("multiplySquare",
      checkForAll Gen.square Multiply.multiplySquare) :
   ("squareSquare",
      checkForAll Gen.square Multiply.squareSquare) :
   ("power",
      Gen.withExtra checkForAll
         (QC.choose (0,10)) Gen.square Multiply.power) :
   ("congruence",
      checkForAll ((,) <$> Gen.square <#*#> Gen.matrix) congruence) :
   ("congruenceAdjoint",
      checkForAll ((,) <$> Gen.matrix <#*#> Gen.square) congruenceAdjoint) :
   ("multiplySquare",
      checkForAll Gen.square multiplySquare) :
   ("multiplyPower",
      Gen.withExtra checkForAll (QC.choose (0,10)) Gen.square multiplyPower) :
   ("multiplyInverse",
      checkForAll Gen.invertible multiplyInverse) :
   ("determinantSingleton",
      checkForAll Gen.scalar determinantSingleton) :
   ("determinantTranspose",
      checkForAll Gen.square determinantTranspose) :
   ("multiplyDeterminant",
      checkForAll ((,) <$> Gen.square <#*#> Gen.square) multiplyDeterminant) :
   ("multiplySolve",
      checkForAll ((,) <$> Gen.invertible <#\#> Gen.matrix) multiplySolve) :
   Divide.testsVar Gen.invertible ++

   ("schur",
      checkForAll Gen.square schur) :
   ("schurComplex",
      checkForAll (Matrix.toComplex <$> Gen.square) schurComplex) :
   ("eigensystem",
      checkForAll genDiagonalizable eigensystem) :
   ("eigensystemLeft",
      checkForAll genDiagonalizable eigensystemLeft) :
   ("eigensystemRight",
      checkForAll genDiagonalizable eigensystemRight) :
   []