{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE GADTs #-}
module Test.Triangular (
   testsVar,
   genTriangular,
   PowerStrip(..),
   diagonal,
   ) where

import qualified Test.Mosaic as Mosaic
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.Logic as Logic
import qualified Test.Utility as Util
import Test.Mosaic (repack)
import Test.Generator ((<#*|>), (<#*#>), (<#\#>))
import Test.Utility
         (approx, approxArray, approxMatrix,
          equalArray, Tagged, (!|||), (!===))

import qualified Numeric.LAPACK.Matrix.Quadratic as Quadratic
import qualified Numeric.LAPACK.Matrix.Triangular as Triangular
import qualified Numeric.LAPACK.Matrix.Diagonal as Diagonal
import qualified Numeric.LAPACK.Matrix.Shape as MatrixShape
import qualified Numeric.LAPACK.Matrix.Array as ArrMatrix
import qualified Numeric.LAPACK.Matrix.Layout as Layout
import qualified Numeric.LAPACK.Matrix as Matrix
import qualified Numeric.LAPACK.Vector as Vector
import Numeric.LAPACK.Matrix.Triangular (Triangular)
import Numeric.LAPACK.Matrix (General, ShapeInt, (|||), (===))
import Numeric.LAPACK.Vector (Vector)
import Numeric.LAPACK.Scalar (RealOf, selectReal)

import qualified Numeric.Netlib.Class as Class

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

import Control.Applicative ((<$>))

import qualified Data.List as List
import Data.Tuple.HT (uncurry3)
import Data.Semigroup ((<>))

import qualified Test.QuickCheck as QC


data PowerStrip lower upper where
   Diagonal :: PowerStrip MatrixShape.Empty MatrixShape.Empty
   Lower :: PowerStrip MatrixShape.Filled MatrixShape.Empty
   Upper :: PowerStrip MatrixShape.Empty MatrixShape.Filled

genTriangular ::
   (MatrixShape.TriDiag diag, MatrixShape.DiagUpLo lo up) =>
   (Logic.Dim sh, Shape.Indexed sh, Shape.Index sh ~ ix, Eq ix,
    Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   PowerStrip lo up ->
   MatrixShape.DiagSingleton diag ->
   Layout.PackingSingleton pack ->
   Gen.Square sh a (ArrMatrix.Quadratic pack diag lo up sh a)
genTriangular _ _ p = repack p <$> Gen.triangular


type TriangularP pack lo diag up sh = ArrMatrix.Quadratic pack diag lo up sh
type FlexDiagonalP pack diag sh =
      ArrMatrix.Quadratic pack diag MatrixShape.Empty MatrixShape.Empty sh
type FlexLowerP pack diag sh =
      ArrMatrix.Quadratic pack diag MatrixShape.Filled MatrixShape.Empty sh
type FlexUpperP pack diag sh =
      ArrMatrix.Quadratic pack diag MatrixShape.Empty MatrixShape.Filled sh

expandTriangle ::
   (MatrixShape.PowerStrip lo, MatrixShape.TriDiag diag, MatrixShape.PowerStrip up,
    Shape.C sh, Class.Floating a) =>
   TriangularP pack lo diag up sh a -> General sh sh a
expandTriangle = Matrix.fromFull . Matrix.toFull

transposedZero ::
   (Shape.C height, Shape.C width, Class.Floating a) =>
   General height width a -> General width height a
transposedZero = ArrMatrix.zero . ArrMatrix.shape . Matrix.transpose

stackDiagonal ::
   (Layout.Packing pack) =>
   (MatrixShape.TriDiag diag, Class.Floating a) =>
   (FlexDiagonalP pack diag ShapeInt a,
    FlexDiagonalP pack diag ShapeInt a) ->
   Bool
stackDiagonal (a,c) =
   let b = Matrix.zero $
           MatrixShape.general MatrixShape.RowMajor
              (Matrix.height a) (Matrix.height c)
   in equalArray
         (expandTriangle $ Diagonal.stack a c)
         (expandTriangle a ||| b
          ===
          Matrix.transpose b ||| expandTriangle c)

stackLower ::
   (Layout.Packing pack) =>
   (MatrixShape.TriDiag diag, Class.Floating a) =>
   (FlexLowerP pack diag ShapeInt a,
    General ShapeInt ShapeInt a,
    FlexLowerP pack diag ShapeInt a) ->
   Bool
stackLower (a,b,c) =
   equalArray
      (expandTriangle $ Triangular.stackLower a b c)
      (expandTriangle a !||| transposedZero b
       ===
       b !||| expandTriangle c)

stackUpper ::
   (Layout.Packing pack) =>
   (MatrixShape.TriDiag diag, Class.Floating a) =>
   (FlexUpperP pack diag ShapeInt a,
    General ShapeInt ShapeInt a,
    FlexUpperP pack diag ShapeInt a) ->
   Bool
stackUpper (a,b,c) =
   equalArray
      (expandTriangle $ Triangular.stackUpper a b c)
      (expandTriangle a ||| b
       !===
       transposedZero b ||| expandTriangle c)


splitDiagonal ::
   (Layout.Packing pack) =>
   (MatrixShape.TriDiag diag, Class.Floating a) =>
   FlexDiagonalP pack diag (ShapeInt::+ShapeInt) a -> Bool
splitDiagonal abc =
   equalArray abc $
      uncurry Diagonal.stack $ Diagonal.split abc

splitLower ::
   (Layout.Packing pack) =>
   (MatrixShape.TriDiag diag, Class.Floating a) =>
   FlexLowerP pack diag (ShapeInt::+ShapeInt) a -> Bool
splitLower abc =
   equalArray abc $
      uncurry3 Triangular.stackLower $ Triangular.splitLower abc

splitUpper ::
   (Layout.Packing pack) =>
   (MatrixShape.TriDiag diag, Class.Floating a) =>
   FlexUpperP pack diag (ShapeInt::+ShapeInt) a -> Bool
splitUpper abc =
   equalArray abc $
      uncurry3 Triangular.stackUpper $ Triangular.splitUpper abc


multiplyIdentity ::
   (Layout.Packing pack) =>
   (MatrixShape.DiagUpLo lo up, MatrixShape.TriDiag diag,
    Eq lo, Eq up,
    Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (TriangularP pack lo diag up ShapeInt a,
    TriangularP pack lo diag up ShapeInt a) ->
   Bool
multiplyIdentity (eye,a) = approxArray a (eye <> a)

multiply ::
   (Layout.Packing pack) =>
   (MatrixShape.DiagUpLo lo up, MatrixShape.TriDiag diag,
    Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   (TriangularP pack lo diag up ShapeInt a,
    TriangularP pack lo diag up ShapeInt a) ->
   Bool
multiply (a,b) =
   approxArray
      (Matrix.toFull a <> Matrix.toFull b)
      (Matrix.toFull $ a <> b)



genInvertible ::
   (MatrixShape.DiagUpLo lo up, MatrixShape.TriDiag diag,
    Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   PowerStrip lo up ->
   MatrixShape.DiagSingleton diag ->
   Layout.PackingSingleton pack ->
   GenTriangularP pack lo diag up ShapeInt a
genInvertible _cont _diag pack =
   repack pack <$> Gen.condition Util.invertible Gen.triangular



eigenvaluesDeterminant ::
   (Layout.Packing pack) =>
   (MatrixShape.DiagUpLo lo up, MatrixShape.TriDiag diag,
    Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   TriangularP pack lo diag up ShapeInt a -> Bool
eigenvaluesDeterminant a =
   approx
      (selectReal 1e-1 1e-5)
      (Matrix.determinant a)
      (Vector.product $ Triangular.eigenvalues a)


genDiagonalizable ::
   (MatrixShape.DiagUpLo lo up,
    Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   PowerStrip lo up ->
   Layout.PackingSingleton pack ->
   GenTriangularP pack lo MatrixShape.Arbitrary up ShapeInt a
genDiagonalizable _cont pack =
   fmap (repack pack) $
   flip Gen.mapGen Gen.triangularShape $ \maxElem shape -> do
      d <- Util.genDistinct [-3..3] [-10..10] $ MatrixShape.squareSize shape
      Util.genArrayExtraDiag maxElem shape (\r -> return (d!r))

diagonal ::
   (Layout.Packing pack) =>
   (MatrixShape.DiagUpLo lo up, Shape.C sh, Class.Floating a) =>
   MatrixShape.Order ->
   Vector sh a -> TriangularP pack lo MatrixShape.Arbitrary up sh a
diagonal order v =
   repack Layout.autoPacking $
   getDiagonal $
   MatrixShape.switchDiagUpLo
      (Diagonal_ $ Quadratic.diagonal order v)
      (Diagonal_ $ Quadratic.diagonal order v)
      (Diagonal_ $ Quadratic.diagonal order v)

newtype Diagonal_ sh a lo up =
   Diagonal_ {getDiagonal :: Triangular lo MatrixShape.Arbitrary up sh a}

eigensystem ::
   (Layout.Packing pack) =>
   (MatrixShape.DiagUpLo lo up, Eq lo, Eq up,
    Class.Floating a, RealOf a ~ ar, Class.Real ar) =>
   MatrixShape.Order ->
   TriangularP pack lo MatrixShape.Arbitrary up ShapeInt a -> Bool
eigensystem order a =
   let (vr,d,vl) = Triangular.eigensystem a
       scal = Triangular.takeDiagonal $ vl <> vr
   in approxMatrix (selectReal 1e-3 1e-5) a
         (vr
          <>
          diagonal order (Vector.divide d scal)
          <>
          vl)


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


type GenTriangularP pack lo diag up sh a =
      Gen.Square sh a (TriangularP pack lo diag up sh a)


propCont ::
   (MatrixShape.TriDiag diag) =>
   MatrixShape.DiagSingleton diag ->
   PowerStrip lo up -> Mosaic.Property diag lo up
propCont _diag cont =
   case cont of
      Diagonal -> Mosaic.Diagonal
      Lower -> Mosaic.Lower
      Upper -> Mosaic.Upper

testsVar ::
   (Show a, Show ar, Class.Floating a, Eq a, RealOf a ~ ar, Class.Real ar) =>
   [(String, Tagged a QC.Property)]
testsVar =
   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 =
   Util.prefix "Diagonal" (testsVarPowerStrip Diagonal p) ++
   Util.prefix "Lower" (testsVarPowerStrip Lower p) ++
   Util.prefix "Upper" (testsVarPowerStrip Upper p) ++
   []

testsVarPowerStrip ::
   (MatrixShape.DiagUpLo lo up, Eq lo, Eq up) =>
   (Layout.Packing pack) =>
   (Show a, Show ar, Class.Floating a, Eq a, RealOf a ~ ar, Class.Real ar) =>
   PowerStrip lo up ->
   Layout.PackingSingleton pack -> [(String, Tagged a QC.Property)]
testsVarPowerStrip cont p =
   Util.prefix "Unit" (testsVarExt cont MatrixShape.Unit p) ++
   Util.prefix "Arbitrary" (testsVarExt cont MatrixShape.Arbitrary p) ++

   Generic.testsDistributive
      (Gen.asMatrixInt $ genTriangular cont MatrixShape.Arbitrary p) ++

   ("eigensystem",
      checkForAllExtra Util.genOrder
         (genDiagonalizable cont p) eigensystem) :
   []

testsVarExt ::
   (Layout.Packing pack) =>
   (MatrixShape.DiagUpLo lo up, Eq lo, Eq up) =>
   (MatrixShape.TriDiag diag) =>
   (Show a, Show ar, Class.Floating a, Eq a, RealOf a ~ ar, Class.Real ar) =>
   PowerStrip lo up ->
   MatrixShape.DiagSingleton diag ->
   Layout.PackingSingleton pack ->
   [(String, Tagged a QC.Property)]
testsVarExt cont diag p =
   ("index",
      checkForAll
         (Indexed.genMatrixIndex (genTriangular cont diag p))
         Indexed.unitDot) :
   ("stack",
      case cont of
         Diagonal ->
            let gen = genTriangular cont diag p
            in checkForAll (Gen.stackDiagonal gen gen) stackDiagonal
         Lower ->
            let gen = genTriangular cont diag p
            in checkForAll
                  (Gen.stack3 gen (Matrix.transpose <$> Gen.matrix) gen)
                  stackLower
         Upper ->
            let gen = genTriangular cont diag p
            in checkForAll (Gen.stack3 gen Gen.matrix gen) stackUpper) :
   ("split",
      case cont of
         Diagonal -> checkForAll (genTriangular cont diag p) splitDiagonal
         Lower -> checkForAll (genTriangular cont diag p) splitLower
         Upper -> checkForAll (genTriangular cont diag p) splitUpper) :
   ("forceOrder",
      checkForAllExtra Util.genOrder
         ((,) <$> genTriangular cont diag p <#*|> Gen.vector)
         Generic.forceOrder) :
   ("forceOrderInverse",
      checkForAll (genTriangular cont diag p) Generic.forceOrderInverse) :

   Mosaic.testsVar (propCont diag cont) p ++

   ("multiplyIdentity",
      checkForAll
         ((,) <$> Mosaic.genIdentity (propCont diag cont) p
              <#*#> genTriangular cont diag p)
         multiplyIdentity) :
   ("multiply",
      let gen = genTriangular cont diag p
      in checkForAll ((,) <$> gen <#*#> gen) multiply) :

   ("determinant",
      checkForAll (genTriangular cont diag p) Divide.determinant) :
   ("solve",
      checkForAll
         ((,) <$> genInvertible cont diag p <#\#> Gen.matrix)
         Divide.solve) :
   ("solveIdentity",
      checkForAll
         ((,) <$> Mosaic.genIdentity (propCont diag cont) p <#\#> Gen.matrix)
         Divide.solveIdentity) :
   ("inverse",
      checkForAll (genInvertible cont diag p) Divide.inverse) :
   map
      (\(name,test) ->
         (name, test $ fmap Divide.SquareMatrix $ genInvertible cont diag p))
      Divide.testsVarAny ++

   ("eigenvaluesDeterminant",
      checkForAll (genTriangular cont diag p) eigenvaluesDeterminant) :
   []