{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleContexts #-}
module Test.Banded (testsVar) where

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.Utility as Util
import Test.Banded.Utility
         (Square(Square), genSquare, shapeBandedFromFull,
          offDiagonals, offDiagonalNats)
import Test.Generator ((<#=#>), (<-*#>), (<#*|>), (<-*|>), (<#*#>), (<#\#>))
import Test.Logic (Dim)
import Test.Utility
         (approx, approxArray, approxMatrix, approxVector,
          genOrder, genArray, genArrayIndexed,
          equalArray, Tagged, equalListWith)

import qualified Numeric.LAPACK.Matrix.Banded.Naive as BandedNaive
import qualified Numeric.LAPACK.Matrix.Banded as Banded
import qualified Numeric.LAPACK.Matrix.Shape as MatrixShape
import qualified Numeric.LAPACK.Matrix.Shape.Omni as Omni
import qualified Numeric.LAPACK.Matrix.Layout as Layout
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 (ShapeInt, (#*##), (-*#), (#*|))
import Numeric.LAPACK.Vector (Vector)
import Numeric.LAPACK.Scalar (RealOf, absolute, one)

import qualified Numeric.Netlib.Class as Class

import qualified Type.Data.Num.Unary.Literal as TypeNum
import qualified Type.Data.Num.Unary.Proof as Proof
import qualified Type.Data.Num.Unary as Unary
import Type.Data.Num.Unary (unary, (:+:))

import qualified Data.Array.Comfort.Storable as Array
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Shape ((::+))

import Foreign.Storable (Storable)

import Control.Applicative (liftA2, (<$>))

import Data.Tuple.HT (mapPair, mapSnd)

import qualified Test.QuickCheck as QC


data Banded height width a =
   forall sub super.
   (Unary.Natural sub, Unary.Natural super) =>
   Banded (Banded.General sub super height width a)

instance
   (Show width, Show height, Show a,
    Shape.C width, Shape.C height, Storable a) =>
      Show (Banded height width a) where
   showsPrec p (Banded a) = showsPrec p a

genBanded ::
   (Dim height, Dim width, Class.Floating a) =>
   Gen.Matrix height width a (Banded height width a)
genBanded = flip Gen.mapGenDim Gen.matrixShape $ \maxElem maxDim shape -> do
   kl <- QC.choose (0, toInteger maxDim)
   ku <- QC.choose (0, toInteger maxDim)
   Unary.reifyNatural kl $ \sub ->
      Unary.reifyNatural ku $ \super ->
      fmap Banded $ genArray maxElem $
      shapeBandedFromFull (unary sub, unary super) shape

genBandedLimitted ::
   (Dim height, Dim width, Shape.Indexed height, Shape.Indexed width,
    Class.Floating a) =>
   Gen.Matrix height width a (Banded height width a)
genBandedLimitted =
      flip Gen.mapGenDim Gen.matrixShape $ \maxElem maxDim shape -> do
   kl <- QC.choose (0, toInteger maxDim)
   ku <- QC.choose (0, toInteger maxDim)
   Unary.reifyNatural kl $ \sub ->
      Unary.reifyNatural ku $ \super ->
      fmap (Banded . ArrMatrix.lift0) $
      Util.genArrayIndexed
            (Omni.toBanded $
             shapeBandedFromFull (unary sub, unary super) shape) $ \ix ->
         case ix of
            Layout.InsideBox _ _ -> Util.genElement maxElem
            Layout.VertOutsideBox _ _ -> return 0
            Layout.HorizOutsideBox _ _ -> return 0


data Banded2 height width a =
   forall sub super.
   (Unary.Natural sub, Unary.Natural super) =>
   Banded2
      (Banded.General sub super height width a)
      (Banded.General sub super height width a)

instance
   (Show width, Show height, Show a,
    Shape.C width, Shape.C height, Storable a) =>
      Show (Banded2 height width a) where
   showsPrec p (Banded2 a b) =
      showParen True $ showsPrec p a . showString ", " . showsPrec p b

genMatrixShape ::
   (Dim height, Dim width) =>
   Gen.Matrix height width a (MatrixShape.General height width)
genMatrixShape = Gen.mapGen (const return) Gen.matrixShape

genBanded2 ::
   (Dim height, Eq height, Dim width, Eq width, Class.Floating a) =>
   Gen.Matrix height width a (Banded2 height width a)
genBanded2 =
   flip Gen.mapQCDim ((,) <$> genMatrixShape <#=#> genMatrixShape) $
      \maxElem maxDim (shapeA,shapeB) -> do
   kl <- QC.choose (0, toInteger maxDim)
   ku <- QC.choose (0, toInteger maxDim)
   Unary.reifyNatural kl $ \sub ->
      Unary.reifyNatural ku $ \super ->
      liftA2 Banded2
         (genArray maxElem $
          shapeBandedFromFull (unary sub, unary super) shapeA)
         (genArray maxElem $
          shapeBandedFromFull (unary sub, unary super) shapeB)


toFull ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Banded ShapeInt ShapeInt a -> Bool
toFull (Banded a) = equalArray (Banded.toFull a) (BandedNaive.toFull a)

forceOrderIndexed ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Layout.Order -> Banded ShapeInt ShapeInt a -> Bool
forceOrderIndexed newOrder (Banded a) =
   equalArray (Banded.forceOrder newOrder a) (BandedNaive.forceOrder newOrder a)


takeTopLeftSquare ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Square (ShapeInt ::+ ShapeInt) a -> Bool
takeTopLeftSquare (Square a) =
   approxArray
      (Square.takeTopLeft $ Banded.toFull a)
      (Banded.toFull $ Banded.takeTopLeftSquare a)

takeBottomRightSquare ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Square (ShapeInt ::+ ShapeInt) a -> Bool
takeBottomRightSquare (Square a) =
   approxArray
      (Square.takeBottomRight $ Banded.toFull a)
      (Banded.toFull $ Banded.takeBottomRightSquare a)


multiplyFullIdentity ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Banded ShapeInt ShapeInt a -> Bool
multiplyFullIdentity (Banded m) =
   let a = Banded.toFull m
   in approxArray a $
      Banded.multiplyFull m $ Square.toGeneral $ Square.identityFromWidth a


multiplyVectorDot ::
   (Class.Floating a, Eq a) =>
   (Vector ShapeInt a,
    Banded ShapeInt ShapeInt a,
    Vector ShapeInt a) ->
   Bool
multiplyVectorDot (x, Banded m, y) =
   Vector.dot x (m#*|y) == Vector.dot (x-*#m) y


multiplyFullAny ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Banded ShapeInt ShapeInt a,
    Matrix.General ShapeInt ShapeInt a) ->
   Bool
multiplyFullAny (Banded a, b) =
   approxArray
      (Banded.multiplyFull a b)
      (Matrix.multiply (Banded.toFull a) b)

multiplyFullColumns ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Banded ShapeInt ShapeInt a,
    Matrix.General ShapeInt ShapeInt a) ->
   Bool
multiplyFullColumns (Banded a, b) =
   equalListWith approxVector
      (Matrix.toColumns (Banded.multiplyFull a b))
      (map (Banded.multiplyVector a) (Matrix.toColumns b))


multiplyFullAssoc ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Banded ShapeInt ShapeInt a,
    Matrix.General ShapeInt ShapeInt a,
    Matrix.General ShapeInt ShapeInt a) ->
   Bool
multiplyFullAssoc (Banded a, b, c) =
   approxArray
      (Matrix.multiply (Banded.multiplyFull a b) c)
      (Banded.multiplyFull a (Matrix.multiply b c))


addOffDiagonals ::
   (Unary.Natural subA, Unary.Natural superA,
    Unary.Natural subB, Unary.Natural superB) =>
   Banded.General subA superA heightA widthA a ->
   Banded.General subB superB heightB widthB a ->
   (Proof.Nat (subA :+: subB), Proof.Nat (superA :+: superB))
addOffDiagonals a b =
   fst $ Layout.addOffDiagonals (offDiagonals a) (offDiagonals b)

multiplyBanded ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Banded ShapeInt ShapeInt a,
    Banded ShapeInt ShapeInt a) ->
   Bool
multiplyBanded (Banded a, Banded b) =
   case addOffDiagonals a b of
      (Proof.Nat, Proof.Nat) ->
         approxArray
            (Banded.toFull (Banded.multiply a b))
            (Banded.multiplyFull a (Banded.toFull b))

multiplyBandedVectorAssoc ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Banded ShapeInt ShapeInt a,
    Banded ShapeInt ShapeInt a,
    Vector ShapeInt a) ->
   Bool
multiplyBandedVectorAssoc (Banded a, Banded b, x) =
   case addOffDiagonals a b of
      (Proof.Nat, Proof.Nat) ->
         approxVector (a #*| b #*| x) (Banded.multiply a b #*| x)


multiplyBandedAssoc ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Banded ShapeInt ShapeInt a,
    Banded ShapeInt ShapeInt a,
    Banded ShapeInt ShapeInt a) ->
   Bool
multiplyBandedAssoc (Banded a, Banded b, Banded c) =
   let ab = Banded.multiply a b
       bc = Banded.multiply b c
       (subA,superA) = offDiagonalNats a
       (subB,superB) = offDiagonalNats b
       (subC,superC) = offDiagonalNats c
   in case (addOffDiagonals a b, addOffDiagonals b c) of
         ((Proof.Nat, Proof.Nat), (Proof.Nat, Proof.Nat)) ->
            case ((addOffDiagonals ab c, addOffDiagonals a bc),
                  (Proof.addAssoc subA subB subC,
                   Proof.addAssoc superA superB superC)) of
               (((Proof.Nat, Proof.Nat), (Proof.Nat, Proof.Nat)),
                (Proof.AddAssoc, Proof.AddAssoc)) ->
                  approxArray (Banded.multiply a bc) (Banded.multiply ab c)


data Upper size a =
   forall super. (Unary.Natural super) => Upper (Banded.Upper super size a)

instance
   (Show size, Show a, Shape.C size, Storable a) =>
      Show (Upper size a) where
   showsPrec p (Upper a) = showsPrec p a

genUpper :: (Class.Floating a) => Gen.MatrixInt a (Upper ShapeInt a)
genUpper = flip Gen.mapGenDim Gen.squareShape $ \maxElem maxDim shape -> do
   ku <- QC.choose (0, toInteger maxDim)
   Unary.reifyNatural ku $ \super ->
      fmap Upper $ genArray maxElem $
      shapeBandedFromFull (unary TypeNum.u0, unary super) shape

multiplyUpperVector ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Upper ShapeInt a, Vector ShapeInt a) -> Bool
multiplyUpperVector (Upper m, x) =
   approxVector (m#*|x) (Banded.toUpperTriangular m #*| x)

multiplyLowerVector ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Upper ShapeInt a, Vector ShapeInt a) -> Bool
multiplyLowerVector (Upper up, x) =
   let lo = Banded.transpose up
   in approxVector (lo#*|x) (Banded.toLowerTriangular lo #*| x)


data UnitUpper size a =
   forall super.
      (Unary.Natural super) => UnitUpper (Banded.UnitUpper super size a)

instance
   (Show size, Show a, Shape.C size, Storable a) =>
      Show (UnitUpper size a) where
   showsPrec p (UnitUpper a) = showsPrec p a

genUnitUpper :: (Class.Floating a) => Gen.MatrixInt a (UnitUpper ShapeInt a)
genUnitUpper = flip Gen.mapGenDim Gen.squareShape $ \maxElem maxDim shape -> do
   ku <- QC.choose (0, toInteger maxDim)
   Unary.reifyNatural ku $ \super ->
      fmap UnitUpper $ genArrayUnitDiagFromSquare maxElem $
      shapeBandedFromFull (unary TypeNum.u0, unary super) shape

singletonFromSuper ::
   (Unary.Natural super) =>
   MatrixShape.BandedUpper super sh -> Unary.HeadSingleton super
singletonFromSuper _ = Unary.headSingleton

genArrayUnitDiagFromSquare ::
   (Unary.Natural super,
    Shape.Indexed sh, Shape.Index sh ~ i, Eq i, Class.Floating a) =>
   Integer ->
   MatrixShape.BandedUpper super sh ->
   QC.Gen (Banded.UnitUpper super sh a)
genArrayUnitDiagFromSquare maxElem shape =
   case singletonFromSuper shape of
      Unary.Zero -> genArrayUnitDiag maxElem $ shapeUnitBanded shape
      Unary.Succ -> genArrayUnitDiag maxElem $ shapeUnitBanded shape

genArrayUnitDiag ::
   (Omni.BandedTriangular super Unary.Zero,
    Omni.BandedTriangular Unary.Zero super,
    Shape.Indexed sh, Shape.Index sh ~ i, Eq i, Class.Floating a) =>
   Integer ->
   MatrixShape.BandedUnitUpper super sh ->
   QC.Gen (Banded.UnitUpper super sh a)
genArrayUnitDiag maxElem shape =
   fmap (ArrMatrix.Array . Array.reshape shape) $
   genArrayIndexed (Omni.toPlain shape) $
      \ix ->
         if case ix of Layout.InsideBox r c -> r==c; _ -> False
            then return one
            else Util.genElement maxElem

shapeUnitBanded ::
   (Omni.BandedTriangular super Unary.Zero,
    Omni.BandedTriangular Unary.Zero super) =>
   MatrixShape.BandedUpper super sh -> MatrixShape.BandedUnitUpper super sh
shapeUnitBanded (Omni.Banded shape) = Omni.UnitBandedTriangular shape



determinant ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Square ShapeInt a -> Bool
determinant (Square a) =
   approx 0.5 (Banded.determinant a) (Square.determinant $ Banded.toFull a)


invertible ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   Square ShapeInt a -> Bool
invertible (Square a) = absolute (Banded.determinant a) > 0.1

multiplySolve ::
   (Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (Square ShapeInt a, Matrix.General ShapeInt ShapeInt a) -> Bool
multiplySolve (Square a, b) =
   approxMatrix 1e-2 (a #*## Banded.solve a b) 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 10 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 =
   ("index",
      checkForAll
         (Indexed.genMatrixIndexGen
            (\(Banded arr) -> Matrix.indices arr)
            genBanded)
         (\(mix, Banded arr) -> Indexed.unitDot (mix, arr))) :
   ("forceOrder",
      checkForAllExtra genOrder
         ((,) <$> Gen.asMatrixInt genBanded <#*|> Gen.vector)
         (\order (Banded a, v) -> Generic.forceOrder order (a,v))) :
   ("forceOrderInverse",
      checkForAll (Gen.asMatrixInt genBandedLimitted)
         (\(Banded a) -> Generic.forceOrderInverse a)) :
   ("forceOrderIndexed",
      checkForAllExtra genOrder genBanded forceOrderIndexed) :
   ("addDistributive",
      checkForAll
         (Generic.genDistribution2 (Gen.asMatrixInt genBanded2))
         (\(Banded2 a b, v) -> Generic.addDistributive ((a,b),v))) :
   ("subDistributive",
      checkForAll
         (Generic.genDistribution2 (Gen.asMatrixInt genBanded2))
         (\(Banded2 a b, v) -> Generic.subDistributive ((a,b),v))) :

   ("toFull",
      checkForAll genBanded toFull) :
   ("takeTopLeftSquare",
      checkForAll genSquare takeTopLeftSquare) :
   ("takeBottomRightSquare",
      checkForAll genSquare takeBottomRightSquare) :

   ("multiplyFullIdentity",
      checkForAll genBanded multiplyFullIdentity) :
   ("multiplyFullAny",
      checkForAll ((,) <$> genBanded <#*#> Gen.matrix) multiplyFullAny) :
   ("multiplyVectorDot",
      checkForAll
         ((,,) <$> Gen.vector <-*#> genBanded <-*|> Gen.vector)
         multiplyVectorDot) :
   ("multiplyFullColumns",
      checkForAll ((,) <$> genBanded <#*#> Gen.matrix) multiplyFullColumns) :
   ("multiplyFullAssoc",
      checkForAll
         ((,,) <$> genBanded <#*#> Gen.matrix <#*#> Gen.matrix)
         multiplyFullAssoc) :
   ("multiplyBanded",
      checkForAll ((,) <$> genBanded <#*#> genBanded) multiplyBanded) :
   ("multiplyBandedVectorAssoc",
      checkForAll
         ((,,) <$> genBanded <#*#> genBanded <#*|> Gen.vector)
         multiplyBandedVectorAssoc) :
   ("multiplyBandedAssoc",
      checkForAll
         ((,,) <$> genBanded <#*#> genBanded <#*#> genBanded)
         multiplyBandedAssoc) :
   ("multiplyUpperVector",
      checkForAll ((,) <$> genUpper <#*|> Gen.vector) multiplyUpperVector) :
   ("multiplyLowerVector",
      checkForAll ((,) <$> genUpper <#*|> Gen.vector) multiplyLowerVector) :
   ("determinant",
      checkForAll genSquare determinant) :
   ("multiplySolve",
      checkForAll
         ((,) <$> Gen.condition invertible genSquare <#\#> Gen.matrix)
         multiplySolve) :
   map
      (mapPair
         (("UnitUpper."++),
          ($ ((\(UnitUpper m) -> Divide.SquareMatrix m) <$> genUnitUpper))))
      Divide.testsVarAny ++
   map
      (mapPair
         (("UnitLower."++),
          ($ ((\(UnitUpper m) -> Divide.SquareMatrix $ Banded.transpose m)
               <$> genUnitUpper))))
      Divide.testsVarAny ++
   map
      (mapSnd
         ($ ((\(Square m) -> Divide.SquareMatrix m)
               <$> Gen.condition invertible genSquare)))
      Divide.testsVarAny ++
   []