{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ConstraintKinds #-}
module Test.Multiply (
   multiplySquare,
   squareSquare,
   power,

   multiplyIdentityVector,
   multiplyIdentityFull,
   multiplyVector,
   multiplyFull,
   multiplyVectorLeft,
   multiplyVectorRight,
   multiplyLeft,
   multiplyRight,

   testsGeneralVar,
   testsVar,
   ) 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, approxArray, approxMatrix, approxVector)

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.Extent as Extent
import qualified Numeric.LAPACK.Matrix as Matrix
import qualified Numeric.LAPACK.Vector as Vector
import Numeric.LAPACK.Matrix (Matrix, 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 Control.Applicative ((<$>))

import qualified Test.QuickCheck as QC



multiplySquare ::
   (Matrix.BoxExtra typ xl, Matrix.BoxExtra typ xu,
    Matrix.Power typ, Matrix.PowerExtra typ xl, Matrix.PowerExtra typ xu,
    Matrix.MultiplySquare typ,
    Matrix.MultiplySquareExtra typ xl, Matrix.MultiplySquareExtra typ xu,
    Matrix.Unpack typ, Matrix.ToQuadratic typ,
    Matrix.UnpackExtra typ xl, Matrix.UnpackExtra typ xu,
    MatrixShape.PowerStrip lower, MatrixShape.PowerStrip upper,
    Shape.C sh, Eq sh, Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Matrix.Quadratic typ xl xu lower upper sh a -> Bool
multiplySquare a =
   Util.approxArray -- (Scalar.selectReal 1e-1 1e-5)
      (Matrix.toFull $ Matrix.square a)
      (a #*## Matrix.toFull a)

squareSquare ::
   (Matrix.Power typ, Matrix.PowerExtra typ xl, Matrix.PowerExtra typ xu,
    Matrix.MultiplySquare typ,
    Matrix.MultiplySquareExtra typ xl, Matrix.MultiplySquareExtra typ xu,
    Matrix.Unpack typ, Matrix.ToQuadratic typ,
    Matrix.UnpackExtra typ xl, Matrix.UnpackExtra typ xu,
    MatrixShape.PowerStrip lower, MatrixShape.PowerStrip upper,
    Shape.C sh, Eq sh, Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Matrix.Quadratic typ xl xu lower upper sh a -> Bool
squareSquare a =
   Util.approxArray -- (Scalar.selectReal 1e-1 1e-5)
      (Matrix.toFull $ Matrix.square a)
      (Square.square $ Matrix.toFull a)

power ::
   (Matrix.BoxExtra typ xl, Matrix.BoxExtra typ xu,
    Matrix.Power typ, Matrix.PowerExtra typ xl, Matrix.PowerExtra typ xu,
    Matrix.MultiplySquare typ,
    Matrix.MultiplySquareExtra typ xl, Matrix.MultiplySquareExtra typ xu,
    Matrix.Unpack typ, Matrix.ToQuadratic typ,
    Matrix.UnpackExtra typ xl, Matrix.UnpackExtra typ xu,
    MatrixShape.PowerStrip lower, MatrixShape.PowerStrip upper,
    Shape.C sh, Eq sh, Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Int -> Matrix.Quadratic typ xl xu lower upper sh a -> Bool
power n0 a =
   let n = fromIntegral n0
       b = Matrix.toFull (Matrix.power (n+1) a)
       c = a #*## Matrix.toFull (Matrix.power n a)
       normInf1 = Vector.normInf1 . ArrMatrix.toVector
   in Util.approxArrayTol (1e-6 * (normInf1 b + normInf1 c)) b c



multiplyIdentityVector ::
   (Matrix.Unpack typ, Matrix.MultiplyVector typ) =>
   (Matrix.UnpackExtra typ xl, Matrix.UnpackExtra typ xu) =>
   (Matrix.MultiplyVectorExtra typ xl, Matrix.MultiplyVectorExtra typ xu) =>
   (MatrixShape.Strip lo, MatrixShape.Strip up,
    Shape.C sh, Eq sh, Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Matrix.Quadratic typ xl xu lo up sh a, Vector sh a) -> Bool
multiplyIdentityVector (eye,a) = approxVector a (eye #*| a)

multiplyIdentityFull ::
   (Matrix.MultiplySquare typ, Matrix.ToQuadratic typ) =>
   (Matrix.MultiplySquareExtra typ xl, Matrix.MultiplySquareExtra typ xu) =>
   (Matrix.BoxExtra typ xl, Matrix.BoxExtra typ xu) =>
   (MatrixShape.Strip lo, MatrixShape.Strip up,
    Shape.C sh, Eq sh, Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Matrix.Quadratic typ xl xu lo up sh a, Matrix.General sh ShapeInt a) ->
   Bool
multiplyIdentityFull (eye,a) = approxArray a (eye #*## a)

multiplyVector ::
   (Matrix.Unpack typ, Matrix.MultiplyVector typ) =>
   (Matrix.UnpackExtra typ xl, Matrix.UnpackExtra typ xu) =>
   (Matrix.MultiplyVectorExtra typ xl, Matrix.MultiplyVectorExtra typ xu) =>
   (MatrixShape.Strip lo, MatrixShape.Strip up,
    Shape.C sh, Eq sh, Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Matrix.Quadratic typ xl xu lo up sh a, Vector sh a) -> Bool
multiplyVector (a,x) = approxVector (Matrix.toFull a #*| x) (a #*| x)

multiplyFull ::
   (Matrix.ToQuadratic typ) =>
   (Matrix.BoxExtra typ xl, Matrix.BoxExtra typ xu) =>
   (Matrix.Unpack typ, Matrix.MultiplySquare typ) =>
   (Matrix.UnpackExtra typ xl, Matrix.UnpackExtra typ xu) =>
   (Matrix.MultiplySquareExtra typ xl, Matrix.MultiplySquareExtra typ xu) =>
   (MatrixShape.Strip lo, MatrixShape.Strip up,
    Shape.C sh, Eq sh, Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Matrix.Quadratic typ xl xu lo up sh a,
    Matrix.General sh ShapeInt a) ->
   Bool
multiplyFull (a,b) = approxArray (Matrix.toFull a #*## b) (a #*## b)


multiplyVectorLeft ::
   (Matrix.Unpack typ, Matrix.MultiplyVector typ) =>
   (Matrix.UnpackExtra typ xl, Matrix.UnpackExtra typ xu) =>
   (Matrix.MultiplyVectorExtra typ xl, Matrix.MultiplyVectorExtra typ xu) =>
   (MatrixShape.Strip lo, MatrixShape.Strip up,
    Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Shape.C height, Eq height, Shape.C width, Eq width,
    Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Vector height a, Matrix typ xl xu lo up meas vert horiz height width a) ->
   Bool
multiplyVectorLeft (x,a) =
   approxVector (x -*# Matrix.toFull a) (x -*# a)

multiplyVectorRight ::
   (Matrix.Unpack typ, Matrix.MultiplyVector typ) =>
   (Matrix.UnpackExtra typ xl, Matrix.UnpackExtra typ xu) =>
   (Matrix.MultiplyVectorExtra typ xl, Matrix.MultiplyVectorExtra typ xu) =>
   (MatrixShape.Strip lo, MatrixShape.Strip up,
    Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Shape.C height, Eq height, Shape.C width, Eq width,
    Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Matrix typ xl xu lo up meas vert horiz height width a, Vector width a) ->
   Bool
multiplyVectorRight (a,x) =
   approxVector (Matrix.toFull a #*| x) (a #*| x)


multiplyLeft ::
   (Matrix.ToQuadratic typ) =>
   (Matrix.BoxExtra typ xl, Matrix.BoxExtra typ xu) =>
   (Matrix.Unpack typ) =>
   (Matrix.UnpackExtra typ xl, Matrix.UnpackExtra typ xu) =>
   (Matrix.MultiplySquare typ) =>
   (Matrix.MultiplySquareExtra typ xl, Matrix.MultiplySquareExtra typ xu) =>
   (MatrixShape.Strip lo, MatrixShape.Strip up) =>
   (Shape.C sh, Eq sh, Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Matrix.General ShapeInt sh a,
    Matrix.Quadratic typ xl xu lo up sh a) -> Bool
multiplyLeft (a,b) =
   approxMatrix 1e-5 (a ##*# Matrix.toFull b) (a ##*# b)

multiplyRight ::
   (Matrix.ToQuadratic typ) =>
   (Matrix.BoxExtra typ xl, Matrix.BoxExtra typ xu) =>
   (Matrix.Unpack typ) =>
   (Matrix.UnpackExtra typ xl, Matrix.UnpackExtra typ xu) =>
   (Matrix.MultiplySquare typ) =>
   (Matrix.MultiplySquareExtra typ xl, Matrix.MultiplySquareExtra typ xu) =>
   (MatrixShape.Strip lo, MatrixShape.Strip up) =>
   (Shape.C sh, Eq sh, Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Matrix.Quadratic typ xl xu lo up sh a,
    Matrix.General sh ShapeInt a) -> Bool
multiplyRight (a,b) =
   approxArray (Matrix.toFull a #*## b) (a #*## b)



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


testsGeneralVar ::
   (Matrix.Box typ, Matrix.BoxExtra typ xl, Matrix.BoxExtra typ xu,
    Matrix.Unpack typ, Matrix.UnpackExtra typ xl, Matrix.UnpackExtra typ xu,
    Matrix.MultiplyVector typ,
    Matrix.MultiplyVectorExtra typ xl, Matrix.MultiplyVectorExtra typ xu) =>
   (MatrixShape.Strip lo, MatrixShape.Strip up) =>
   (Extent.Measure meas, Extent.C vert, Extent.C horiz) =>
   (Matrix typ xl xu lo up meas vert horiz height width a ~ matrix) =>
   (Show matrix) =>
   (Show height, Eq height, Logic.Dim height,
    Show width, Eq width, Logic.Dim width,
    Show a, Class.Floating a, Eq a, RealOf a ~ ar, Class.Real ar) =>
   Gen.Matrix height width a matrix ->
   [(String, Tagged a QC.Property)]
testsGeneralVar genMatrix =
   ("multiplyVectorLeft",
      checkForAll ((,) <$> Gen.vector <-*#> genMatrix) multiplyVectorLeft) :
   ("multiplyVectorRight",
      checkForAll ((,) <$> genMatrix <#*|> Gen.vector) multiplyVectorRight) :
   []

testsVar ::
   (Matrix.ToQuadratic typ) =>
   (Matrix.Box typ, Matrix.BoxExtra typ xl, Matrix.BoxExtra typ xu,
    Matrix.Unpack typ, Matrix.UnpackExtra typ xl, Matrix.UnpackExtra typ xu,
    Matrix.Power typ, Matrix.PowerExtra typ xl, Matrix.PowerExtra typ xu,
    Matrix.MultiplySquare typ,
    Matrix.MultiplySquareExtra typ xl, Matrix.MultiplySquareExtra typ xu) =>
   (Matrix.MultiplyVector typ,
    Matrix.MultiplyVectorExtra typ xl, Matrix.MultiplyVectorExtra typ xu) =>
   (MatrixShape.PowerStrip lo, MatrixShape.PowerStrip up) =>
   (Matrix.Quadratic typ xl xu lo up sh a ~ matrix) =>
   (Show matrix) =>
   (Show sh, Eq sh, Logic.Dim sh,
    Show a, Class.Floating a, Eq a, RealOf a ~ ar, Class.Real ar) =>
   Gen.Square sh a matrix ->
   Gen.Square sh a matrix ->
   [(String, Tagged a QC.Property)]
testsVar genIdentity genMatrix =
   testsGeneralVar genMatrix ++

   ("multiplyIdentityVector",
      checkForAll
         ((,) <$> genIdentity <#*|> Gen.vector)
         multiplyIdentityVector) :
   ("multiplyIdentityFull",
      checkForAll
         ((,) <$> genIdentity <#*#> Gen.matrix)
         multiplyIdentityFull) :
   ("multiplyVector",
      checkForAll ((,) <$> genMatrix <#*|> Gen.vector)
         multiplyVector) :
   ("multiplyFull",
      checkForAll ((,) <$> genMatrix <#*#> Gen.matrix)
         multiplyFull) :
   ("multiplyLeft",
      checkForAll
         ((,) <$> Gen.matrix <#*#> genMatrix)
         multiplyLeft) :
   ("multiplyRight",
      checkForAll
         ((,) <$> genMatrix <#*#> Gen.matrix)
         multiplyRight) :

   ("multiplySquare",
      checkForAll genMatrix multiplySquare) :
   ("squareSquare",
      checkForAll genMatrix squareSquare) :
   ("power",
      checkForAllExtra (QC.choose (0,10)) genMatrix power) :
   []