{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE ExistentialQuantification #-}
module Test.Divide (
   testsVar,
   testsVarAny,
   SquareMatrix(SquareMatrix),
   determinant,
   solve, solveIdentity, inverse,
   approxRelMatrix, approxRelArray,
   ) where

import qualified Test.Generator as Gen
import qualified Test.Logic as Logic
import qualified Test.Utility as Util
import Test.Generator ((<#\#>), (<#/#>))
import Test.Utility (Tagged, approxMatrix)

import qualified Numeric.LAPACK.Linear.LowerUpper as LU
import qualified Numeric.LAPACK.Matrix.Square as Square
import qualified Numeric.LAPACK.Matrix.Extent as Extent
import qualified Numeric.LAPACK.Matrix.Array as ArrMatrix
import qualified Numeric.LAPACK.Matrix.Special as Special
import qualified Numeric.LAPACK.Matrix.Shape as MatrixShape
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, ArrayMatrix)
import Numeric.LAPACK.Matrix (ShapeInt, (##/#), (##*#), (#*##), (#\##))
import Numeric.LAPACK.Scalar (RealOf)

import qualified Numeric.Netlib.Class as Class

import qualified Data.Array.Comfort.Shape as Shape

import Control.Applicative ((<$>))

import Data.Tuple.HT (mapSnd)

import qualified Test.QuickCheck as QC



determinant ::
   (MatrixShape.Packing pack, MatrixShape.Property property,
    MatrixShape.Strip lower, MatrixShape.Strip upper,
    Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Quadratic pack property lower upper ShapeInt a -> Bool
determinant a =
   let d = Square.determinant $ Matrix.toSquare a
   in Util.approx
         (max (Scalar.selectReal 1 1e-2)
              (Scalar.selectReal 1e-3 1e-5 * Scalar.absolute d))
         d
         (Matrix.determinant a)


approxRelMatrix, approxRelArray ::
   (MatrixShape.PowerStrip lower, MatrixShape.PowerStrip upper) =>
   (Extent.Measure meas, Extent.C vert, Extent.C horiz) =>
   (Shape.C height, Shape.C width, Eq height, Eq width) =>
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   ArrayMatrix pack property lower upper meas vert horiz height width a ->
   ArrayMatrix pack property lower upper meas vert horiz height width a -> Bool
approxRelMatrix a b =
   approxMatrix
      (Scalar.selectReal 0.5 1e-4 *
         (max 1 $ Vector.normInf $ ArrMatrix.unwrap a))
      a b

approxRelArray a b =
   Util.approxArrayTol
      (max (Scalar.selectReal 1e-1 1e-4)
           (Scalar.selectReal 1e-1 1e-5 * Vector.normInf (ArrMatrix.unwrap a)))
      a b



multiplySolveTrans ::
   (Shape.C size, Eq size, Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   LU.Transposition ->
   (SquareMatrix size a, Matrix.General size ShapeInt a) -> Bool
multiplySolveTrans trans (SquareMatrix a, b) =
   approxRelMatrix b $
      Matrix.multiplySquare trans a $ Matrix.solve trans a b

multiplySolveRight ::
   (Shape.C size, Eq size, Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (SquareMatrix size a, Matrix.General size ShapeInt a) -> Bool
multiplySolveRight (SquareMatrix a, b) =
   approxRelMatrix b (a #*## (a #\## b))

multiplySolveLeft ::
   (Shape.C size, Eq size, Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Matrix.General ShapeInt size a, SquareMatrix size a) -> Bool
multiplySolveLeft (b, SquareMatrix a) =
   approxRelMatrix b ((b ##/# a) ##*# a)

multiplyInverseRight ::
   (Shape.C size, Eq size, Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (SquareMatrix size a, Matrix.General size ShapeInt a) -> Bool
multiplyInverseRight (SquareMatrix a, b) =
   approxRelMatrix b (a #*## (Special.Inverse a #*## b))

multiplyInverseLeft ::
   (Shape.C size, Eq size, Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Matrix.General ShapeInt size a, SquareMatrix size a) -> Bool
multiplyInverseLeft (b, SquareMatrix a) =
   approxRelMatrix b ((b ##*# Special.Inverse a) ##*# a)


solve ::
   (MatrixShape.Packing pack, MatrixShape.Property property,
    MatrixShape.Strip lower, MatrixShape.Strip upper) =>
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Quadratic pack property upper lower ShapeInt a,
    Matrix.General ShapeInt ShapeInt a) ->
   Bool
solve (a, b) =
   approxRelArray
      (Square.solve (Matrix.toSquare a) b)
      (Matrix.solveRight a b)

solveIdentity ::
   (MatrixShape.Packing pack, MatrixShape.Property property,
    MatrixShape.Strip lower, MatrixShape.Strip upper) =>
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Quadratic pack property upper lower ShapeInt a,
    Matrix.General ShapeInt ShapeInt a) ->
   Bool
solveIdentity (eye, a) =
   approxMatrix (Scalar.selectReal 1e-3 1e-5) a (Matrix.solveRight eye a)

inverse ::
   (MatrixShape.Packing pack, MatrixShape.Property property,
    MatrixShape.PowerStrip lower, MatrixShape.PowerStrip upper) =>
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Quadratic pack property upper lower ShapeInt a -> Bool
inverse a =
   approxRelArray
      (Square.inverse $ Matrix.toSquare a)
      (Matrix.toSquare $ Matrix.inverse 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)


data SquareMatrix size a =
   forall typ xl xu lower upper matrix.
   (Matrix.BoxExtra typ xl, Matrix.BoxExtra typ xu,
    Matrix.Solve typ, Matrix.SolveExtra typ xl, Matrix.SolveExtra typ xu,
    Matrix.MultiplySquare typ,
    Matrix.MultiplySquareExtra typ xl, Matrix.MultiplySquareExtra typ xu,
    Matrix.ToQuadratic typ,
    MatrixShape.Strip lower, MatrixShape.Strip upper,
    Matrix.Quadratic typ xl xu lower upper size a ~ matrix, Show matrix) =>
   SquareMatrix (Matrix.Quadratic typ xl xu lower upper size a)

instance Show (SquareMatrix size a) where
   show (SquareMatrix m) = show m


testsVarAny ::
   (Logic.Dim sh, Shape.C sh, Show sh, Eq sh) =>
   (Show a, Class.Floating a, Eq a, RealOf a ~ ar, Class.Real ar) =>
   [(String, Gen.Square sh a (SquareMatrix sh a) -> Tagged a QC.Property)]
testsVarAny =
   ("multiplySolveTrans",
      \gen ->
         Gen.withExtra checkForAll QC.arbitraryBoundedEnum
            ((,) <$> gen <#\#> Gen.matrix) multiplySolveTrans) :
   ("multiplySolveRight",
      \gen ->
         checkForAll ((,) <$> gen <#\#> Gen.matrix) multiplySolveRight) :
   ("multiplySolveLeft",
      \gen ->
         checkForAll ((,) <$> Gen.matrix <#/#> gen) multiplySolveLeft) :
   ("multiplyInverseRight",
      \gen ->
         checkForAll ((,) <$> gen <#\#> Gen.matrix) multiplyInverseRight) :
   ("multiplyInverseLeft",
      \gen ->
         checkForAll ((,) <$> Gen.matrix <#/#> gen) multiplyInverseLeft) :
   []

testsVar ::
   (Matrix.BoxExtra typ xl, Matrix.BoxExtra typ xu,
    Matrix.Solve typ, Matrix.SolveExtra typ xl, Matrix.SolveExtra typ xu,
    Matrix.MultiplySquare typ,
    Matrix.MultiplySquareExtra typ xl, Matrix.MultiplySquareExtra typ xu,
    Matrix.ToQuadratic typ,
    MatrixShape.Strip lower, MatrixShape.Strip upper,
    Logic.Dim sh, Shape.C sh, Show sh, Eq sh,
    Matrix.Quadratic typ xl xu lower upper sh a ~ matrix, Show matrix,
    Show a, Class.Floating a, Eq a, RealOf a ~ ar, Class.Real ar) =>
   Gen.Square sh a (Matrix.Quadratic typ xl xu lower upper sh a) ->
   [(String, Tagged a QC.Property)]
testsVar gen =
   map (mapSnd ($ (SquareMatrix <$> gen))) testsVarAny