{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ConstraintKinds  #-}
{-# LANGUAGE GADTs #-}
module Test.Function (testsVar, testsReal) where

import qualified Test.Generator as Gen
import qualified Test.Utility as Util
import Test.Symmetric (genSymmetric)
import Test.Hermitian (genHermitian)
import Test.Triangular
         (genTriangular, diagonal, PowerStrip(Diagonal, Lower, Upper))
import Test.Utility (approx, approxMatrix, Tagged)

import qualified Numeric.LAPACK.Orthogonal as Ortho
import qualified Numeric.LAPACK.Matrix.Function as Fn
import qualified Numeric.LAPACK.Matrix.Hermitian as Hermitian
import qualified Numeric.LAPACK.Matrix.Symmetric as Symmetric
import qualified Numeric.LAPACK.Matrix.Triangular as Triangular
import qualified Numeric.LAPACK.Matrix.Quadratic as Quadratic
import qualified Numeric.LAPACK.Matrix.Square as Square
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 qualified Numeric.LAPACK.Scalar as Scalar
import Numeric.LAPACK.Matrix.Array (Quadratic)
import Numeric.LAPACK.Matrix.Shape (DiagSingleton(Unit, Arbitrary))
import Numeric.LAPACK.Matrix
         (Square, ShapeInt, (#+#), (.*#), (#*##), (\*#), (#\##))
import Numeric.LAPACK.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 Control.Applicative (liftA2, (<$>))

import qualified Data.List as List
import Data.Complex (Complex((:+)))
import Data.Semigroup ((<>))

import qualified Test.QuickCheck as QC



genPositiveSpectrum :: (Class.Real a) => Gen.MatrixInt a (Square ShapeInt a)
genPositiveSpectrum =
   flip fmap Gen.square $ \a ->
      let (q,r) = Ortho.householderTall a
      in q <> diagonalPositive r #*## Matrix.adjoint q

genPositiveHermitian ::
   (Layout.Packing pack, Class.Real a, RealOf a ~ a) =>
   Layout.PackingSingleton pack ->
   Gen.MatrixInt a (HermitianP pack ShapeInt a)
genPositiveHermitian _p =
   flip fmap Gen.square $ \a ->
      let (q,r) = Ortho.householderTall a
          d = Matrix.takeDiagonal r
          dd =
            if Shape.size (Array.shape d) == 0
               then 0
               else 0.1 - min 0 (Vector.minimum d)
      in Hermitian.congruenceDiagonalAdjoint
            (Square.toFull q) (Vector.raise dd d)

genPositiveTriangular ::
   (MatrixShape.DiagUpLo lo up, Class.Real a, RealOf a ~ a) =>
   PowerStrip lo up ->
   MatrixShape.DiagSingleton diag ->
   Layout.PackingSingleton pack ->
   Gen.MatrixInt a (ArrMatrix.Quadratic pack diag lo up ShapeInt a)
genPositiveTriangular cont diag p =
   case diag of
      Unit -> genTriangular cont diag p
      Arbitrary -> fmap diagonalPositive (genTriangular cont diag p)

diagonalPositive ::
   (ArrMatrix.Additive property, ArrMatrix.Scale property, Class.Real a) =>
   (Quadratic pack property lower upper ShapeInt ~ matrix) =>
   matrix a -> matrix a
diagonalPositive a =
   let d =
         if Shape.size (Quadratic.size a) == 0
            then 0
            else 0.1 - min 0 (Vector.minimum (Matrix.takeDiagonal a))
   in a #+# d .*# Matrix.identityFrom a

sqrSqrt ::
   (Fn.SqRt property, MatrixShape.PowerStrip lower, MatrixShape.PowerStrip upper) =>
   (Layout.Packing pack, Class.Real a, RealOf a ~ a) =>
   (Quadratic pack property lower upper ShapeInt ~ matrix) =>
   (matrix a -> matrix a) -> matrix a -> Bool
sqrSqrt sqrtm a =
   approxMatrix (selectReal 1 1e-6) a (Matrix.square (sqrtm a))

sqrtSqr ::
   (Fn.SqRt property, MatrixShape.PowerStrip lower, MatrixShape.PowerStrip upper) =>
   (Layout.Packing pack, Class.Real a, RealOf a ~ a) =>
   (Quadratic pack property lower upper ShapeInt ~ matrix) =>
   (matrix a -> matrix a) -> matrix a -> Bool
sqrtSqr sqrtm a =
   approxMatrix (selectReal 1e-1 1e-6) a (sqrtm (Matrix.square a))

expSum ::
   (Fn.Exp property, ArrMatrix.Additive property) =>
   (MatrixShape.PowerStrip lower, MatrixShape.PowerStrip upper) =>
   (Layout.Packing pack) =>
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Quadratic pack property lower upper ShapeInt ~ matrix) =>
   (b -> matrix a -> matrix a) ->
   (b,b) -> matrix a -> Bool
expSum scale (x,y) m =
   let a = scale x m; b = scale y m
   in approxMatrix (selectReal 10 1e-4)
         (Matrix.toFull $ Fn.exp (a#+#b))
         (Matrix.toFull (Fn.exp a) <> Matrix.toFull (Fn.exp b))

scalarExp :: (Class.Floating a) => a -> a
scalarExp a =
   case Scalar.complexSingletonOf a of
      Scalar.Real -> exp a
      Scalar.Complex -> exp a

expTrace ::
   (Fn.Exp property) =>
   (MatrixShape.PowerStrip lower, MatrixShape.PowerStrip upper) =>
   (Layout.Packing pack) =>
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Quadratic pack property lower upper ShapeInt a -> Bool
expTrace a =
   let e = scalarExp $ Matrix.trace a
   in approx
         (selectReal 1 1e-6 * max 1 (1e-2 * Scalar.absolute e))
         e (Matrix.determinant $ Fn.exp a)

{-
Fails for arbitrary square matrices because eigenvalues can be non-real.
-}
expSqrt ::
   (Fn.Exp property, Fn.SqRt property,
    ArrMatrix.Additive property, ArrMatrix.Homogeneous property) =>
   (MatrixShape.PowerStrip lower, MatrixShape.PowerStrip upper) =>
   (Layout.Packing pack, Class.Real a, RealOf a ~ a) =>
   (Quadratic pack property lower upper ShapeInt ~ matrix) =>
   (matrix a -> matrix a) -> matrix a -> Bool
expSqrt sqrtm a =
   approxMatrix (selectReal 10 1e-4)
      (sqrtm (Fn.exp a))
      (Fn.exp (Matrix.scaleReal (1/2) a))

type HermitianP pack sh =
      ArrMatrix.FullQuadratic pack Omni.HermitianUnknownDefiniteness sh

expSqrtHermitian ::
   (Layout.Packing pack, Class.Real a, RealOf a ~ a) =>
   HermitianP pack ShapeInt a -> Bool
expSqrtHermitian a =
   approxMatrix (selectReal 1e-1 1e-4)
      (Fn.expRealHermitian (Matrix.scaleReal (1/2) a))
      (Fn.sqrt (Fn.expRealHermitian a))

type UnitUpperTriangularP pack sh =
      ArrMatrix.Quadratic pack
         MatrixShape.Unit MatrixShape.Empty MatrixShape.Filled sh

expLogUnipotent ::
   (Layout.Packing pack, Class.Real a, RealOf a ~ a) =>
   UnitUpperTriangularP pack ShapeInt a -> Bool
expLogUnipotent a =
   approxMatrix (selectReal 1e-2 1e-6)
      (Triangular.relaxUnitDiagonal a)
      (Fn.exp (Fn.logUnipotentUpper a))

expLogHermitian ::
   (Layout.Packing pack, Class.Real a, RealOf a ~ a) =>
   HermitianP pack ShapeInt a -> Bool
expLogHermitian a =
   approxMatrix (selectReal 1e-2 1e-6) a (Fn.log (Fn.exp a))


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

genPositiveDistinctTriangular ::
   (Layout.Packing pack, MatrixShape.DiagUpLo lo up) =>
   (Class.Real a, RealOf a ~ ar, Class.Real ar) =>
   (Quadratic pack MatrixShape.Arbitrary lo up ShapeInt ~ matrix) =>
   PowerStrip lo up ->
   Layout.PackingSingleton pack ->
   Gen.MatrixInt a (matrix a)
genPositiveDistinctTriangular cont p =
      flip Gen.mapQC (genTriangular cont Arbitrary p) $ \a -> do
   d <- Util.genDistinct [1..10] [1..10] (Quadratic.size a)
   return $ a #+# diagonal (ArrMatrix.order a) (d |-| Matrix.takeDiagonal a)

expLog ::
   (Fn.Exp property, Fn.Log property) =>
   (MatrixShape.PowerStrip lower, MatrixShape.PowerStrip upper) =>
   (Layout.Packing pack, Class.Real a, RealOf a ~ a) =>
   Quadratic pack property lower upper ShapeInt a -> Bool
expLog a = approxMatrix (selectReal 10 1e-6) a (Fn.exp (Fn.log a))


genCoefficient :: (Class.Floating a) => QC.Gen a
genCoefficient =
   Class.switchFloating
      (QC.choose (-1,1))
      (QC.choose (-1,1))
      (liftA2 (:+) (QC.choose (-1,1)) (QC.choose (-1,1)))
      (liftA2 (:+) (QC.choose (-1,1)) (QC.choose (-1,1)))


checkForAll ::
   (Show a, QC.Testable test) =>
   Gen.T dim tag a -> (a -> test) -> Tagged tag QC.Property
checkForAll gen = Util.checkForAll (Gen.run gen 2 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 =
   ("Square.expSum",
      checkForAllExtra
         (liftA2 (,) genCoefficient genCoefficient)
         Gen.square (expSum (.*#))) :
   ("Square.expTrace",
      checkForAll Gen.square expTrace) :
   concat
      (List.transpose
         [Util.suffix "Packed"   (testsVarPacking Layout.Packed),
          Util.suffix "Unpacked" (testsVarPacking Layout.Unpacked)])

testsVarPacking ::
   (Layout.Packing pack) =>
   (Show a, Show ar, Class.Floating a, Eq a, RealOf a ~ ar, Class.Real ar) =>
   Layout.PackingSingleton pack -> [(String, Tagged a QC.Property)]
testsVarPacking p =
   ("Hermitian.expSum",
      checkForAllExtra
         (liftA2 (,) genCoefficient genCoefficient)
         (genHermitian p) (expSum Matrix.scaleReal)) :
   ("Symmetric.expSum",
      checkForAllExtra
         (liftA2 (,) genCoefficient genCoefficient)
         (genSymmetric p) (expSum (.*#))) :
   ("Diagonal.expSum",
      checkForAllExtra
         (liftA2 (,) genCoefficient genCoefficient)
         (genTriangular Diagonal Arbitrary p) (expSum (.*#))) :
   ("LowerTriangular.expSum",
      checkForAllExtra
         (liftA2 (,) genCoefficient genCoefficient)
         (genTriangular Lower Arbitrary p) (expSum (.*#))) :
   ("UpperTriangular.expSum",
      checkForAllExtra
         (liftA2 (,) genCoefficient genCoefficient)
         (genTriangular Upper Arbitrary p) (expSum (.*#))) :

   ("Hermitian.expTrace",
      checkForAll (genHermitian p) expTrace) :
   ("Symmetric.expTrace",
      checkForAll (genSymmetric p) expTrace) :
   ("Diagonal.expTrace",
      checkForAll (genTriangular Diagonal Arbitrary p) expTrace) :
   ("LowerTriangular.expTrace",
      checkForAll (genTriangular Lower Arbitrary p) expTrace) :
   ("UpperTriangular.expTrace",
      checkForAll (genTriangular Upper Arbitrary p) expTrace) :
   []


testsReal ::
   (Show a, Class.Floating a, Eq a, RealOf a ~ a, Class.Real a) =>
   [(String, Tagged a QC.Property)]
testsReal =
   ("Square.sqrSqrt",
      checkForAll genPositiveSpectrum (sqrSqrt Fn.sqrt)) :
   ("Square.sqrSqrt.DenmanBeavers",
      checkForAll genPositiveSpectrum (sqrSqrt Fn.sqrtDenmanBeavers)) :
{- FixMe:
Schur decomposition seems to produce 2x2 diagonal blocks
even for an entirely real set of eigenvalues.
   ("Square.sqrSqrt.Schur",
      checkForAll genPositiveSpectrum (sqrSqrt Fn.sqrtSchur)) :
-}
   ("expLog",
      checkForAll genPositiveDiagonalizable expLog) :
   concat
      (List.transpose
         [Util.suffix "Packed"   (testsRealPacking Layout.Packed),
          Util.suffix "Unpacked" (testsRealPacking Layout.Unpacked)])

testsRealPacking ::
   (Layout.Packing pack) =>
   (Show a, Class.Floating a, Eq a, RealOf a ~ a, Class.Real a) =>
   Layout.PackingSingleton pack -> [(String, Tagged a QC.Property)]
testsRealPacking p =
   ("Symmetric.sqrSqrt",
      checkForAll (Symmetric.fromHermitian <$> genPositiveHermitian p)
         (sqrSqrt Fn.sqrt)) :
   ("Symmetric.sqrSqrt.DenmanBeavers",
      checkForAll (Symmetric.fromHermitian <$> genPositiveHermitian p)
         (sqrSqrt Fn.sqrtDenmanBeavers)) :
   ("Symmetric.sqrtSqr",
      checkForAll (Symmetric.fromHermitian <$> genPositiveHermitian p)
         (sqrtSqr Fn.sqrt)) :

   ("Hermitian.expSqrt",
      checkForAll (genHermitian p) expSqrtHermitian) :
   ("Symmetric.expSqrt",
      checkForAll (genSymmetric p) (expSqrt Fn.sqrt)) :
   ("Symmetric.expSqrt.DenmanBeavers",
      checkForAll (genSymmetric p)
         (expSqrt Fn.sqrtDenmanBeavers)) :

   ("UpperTriangular.expLogUnipotent",
      checkForAll (genTriangular Upper Unit p) expLogUnipotent) :
   ("UpperTriangular.expLogHermitian",
      checkForAll (genHermitian p) expLogHermitian) :

   Util.prefix "Diagonal" (testsRealPowerStrip Diagonal p) ++
   Util.prefix "Lower" (testsRealPowerStrip Lower p) ++
   Util.prefix "Upper" (testsRealPowerStrip Upper p) ++
   []

testsRealPowerStrip ::
   (Layout.Packing pack) =>
   (MatrixShape.DiagUpLo lo up, Eq lo, Eq up) =>
   (Show a, Class.Floating a, Eq a, RealOf a ~ a, Class.Real a) =>
   PowerStrip lo up ->
   Layout.PackingSingleton pack ->
   [(String, Tagged a QC.Property)]
testsRealPowerStrip cont p =
   ("Unit.sqrSqrt",
      checkForAll (genPositiveTriangular cont Unit p) (sqrSqrt Fn.sqrt)) :
   ("Arbitrary.sqrSqrt",
      checkForAll (genPositiveTriangular cont Arbitrary p) (sqrSqrt Fn.sqrt)) :
   ("sqrSqrt.DenmanBeavers",
      checkForAll (genPositiveTriangular cont Arbitrary p)
         (sqrSqrt Fn.sqrtDenmanBeavers)) :
   ("Unit.sqrtSqr",
      checkForAll (genPositiveTriangular cont Unit p) (sqrtSqr Fn.sqrt)) :
   ("Arbitrary.sqrtSqr",
      checkForAll (genPositiveTriangular cont Arbitrary p) (sqrtSqr Fn.sqrt)) :
   ("expSqrt",
      checkForAll (genTriangular cont Arbitrary p) (expSqrt Fn.sqrt)) :
   ("expSqrt.DenmanBeavers",
      checkForAll (genTriangular cont Arbitrary p)
         (expSqrt Fn.sqrtDenmanBeavers)) :
   ("expLog",
      checkForAll (genPositiveDistinctTriangular cont p) expLog) :
   []