module Test.Permutation where

import qualified Test.Indexed as Indexed
import qualified Test.Generator as Gen
import qualified Test.Utility as Util
import Test.Generator ((<#*#>))
import Test.Logic (Dim)
import Test.Utility (Tagged, equalArray)

import qualified Numeric.LAPACK.Permutation as Perm
import qualified Numeric.LAPACK.Matrix.Permutation as PermMatrix
import qualified Numeric.LAPACK.Matrix.Square as Square
import qualified Numeric.LAPACK.Matrix as Matrix
import qualified Numeric.LAPACK.Vector as Vector
import Numeric.LAPACK.Permutation
         (Permutation, Inversion(Inverted, NonInverted))
import Numeric.LAPACK.Matrix.Square (Square)
import Numeric.LAPACK.Matrix (Matrix, ShapeInt, shapeInt, (#*##))
import Numeric.LAPACK.Vector (Vector)

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.Monad (forM)
import Control.Applicative (liftA2, (<$>))

import Data.Semigroup ((<>))

import qualified Test.QuickCheck as QC


type Pivots = Vector (Perm.Shape ShapeInt) (Perm.Element ShapeInt)

genPivots :: [()] -> QC.Gen Pivots
genPivots nat = do
   let n = length nat
   let nc = fromIntegral n
   fmap (Vector.fromList (Perm.Shape $ shapeInt n)) $
      forM (zip [1..] nat) $ \(i,()) -> Perm.Element <$> QC.choose (i,nc)

genPerm :: [()] -> QC.Gen (Permutation ShapeInt)
genPerm = fmap (Perm.fromPivots NonInverted) . genPivots


permutationPivots :: Inversion -> Pivots -> Bool
permutationPivots inv xs =
   Array.toList (Perm.toPivots inv (Perm.fromPivots inv xs))
   ==
   Array.toList xs

determinantMultiply :: (Permutation ShapeInt, Permutation ShapeInt) -> Bool
determinantMultiply (p0,p1) =
   Perm.determinant (Perm.multiply p0 p1)
   ==
   Perm.determinant p0 <> Perm.determinant p1

transposeMultiply :: (Permutation ShapeInt, Permutation ShapeInt) -> Bool
transposeMultiply (p0,p1) =
   (Array.toList $ Perm.toPivots NonInverted $
    Perm.transpose (Perm.multiply p0 p1))
   ==
   (Array.toList $ Perm.toPivots NonInverted $
    Perm.multiply (Perm.transpose p1) (Perm.transpose p0))


genPermutation :: (Dim sh) => Gen.Matrix sh sh a (Permutation sh)
genPermutation =
   flip Gen.mapGen Gen.squareDim $ \_maxElem sh ->
   let nc = fromIntegral $ Shape.size sh
   in fmap (Perm.fromPivots NonInverted . Vector.fromList (Perm.Shape sh)) $
         forM [1..] $ \i -> Perm.Element <$> QC.choose (i,nc)


permToMatrix ::
   (Shape.C sh, Class.Floating a) =>
   Inversion -> Permutation sh -> Square sh a
permToMatrix inv =
   case inv of
      NonInverted -> Perm.toMatrix
      Inverted -> Matrix.transpose . Perm.toMatrix

applyToMatrix ::
   (Class.Floating a) =>
   Inversion -> (Permutation ShapeInt, Matrix.General ShapeInt ShapeInt a) -> Bool
applyToMatrix inv (p,m) =
   equalArray
      (Perm.apply inv p m)
      (permToMatrix inv p #*## m)

applyMultiply ::
   (Class.Floating a) =>
   Inversion ->
   (Permutation ShapeInt, Permutation ShapeInt,
    Matrix.General ShapeInt ShapeInt a) -> Bool
applyMultiply inv (p0,p1,m) =
   equalArray
      (case inv of
         NonInverted -> Perm.apply inv p0 $ Perm.apply inv p1 m
         Inverted -> Perm.apply inv p1 $ Perm.apply inv p0 m)
      (Perm.apply inv (Perm.multiply p0 p1) m)

applyTranspose ::
   (Class.Floating a) =>
   Inversion -> (Permutation ShapeInt, Matrix.General ShapeInt ShapeInt a) -> Bool
applyTranspose inv (p,m) =
   equalArray
      (Perm.apply inv (Perm.transpose p) m)
      (Matrix.transpose (permToMatrix inv p) #*## m)


genPermMatrix :: (Dim sh) => Gen.Matrix sh sh a (Matrix (Permutation sh) a)
genPermMatrix = PermMatrix.fromPermutation <$> genPermutation

determinantNumber ::
   (Class.Floating a, Eq a) => Matrix (Permutation ShapeInt) a -> Bool
determinantNumber p =
   PermMatrix.determinant p == Square.determinant (PermMatrix.toMatrix p)


tests :: [(String, QC.Property)]
tests =
   ("permutationPivots",
      QC.property $
         QC.forAll QC.arbitraryBoundedEnum $ \inv ->
         QC.forAll (QC.arbitrary >>= genPivots) $ \pivot ->
            permutationPivots inv pivot) :
   ("determinantMultiply",
      QC.property $
         QC.forAll
            (do nat <- QC.arbitrary
                liftA2 (,) (genPerm nat) (genPerm nat))
            determinantMultiply) :
   ("transposeMultiply",
      QC.property $
         QC.forAll
            (do nat <- QC.arbitrary
                liftA2 (,) (genPerm nat) (genPerm nat))
            transposeMultiply) :
   []



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) =>
   [(String, Tagged a QC.Property)]
testsVar =
   ("index",
      checkForAll (Indexed.genMatrixIndex genPermMatrix) Indexed.unitDot) :
   ("applyToMatrix",
      checkForAllExtra QC.arbitraryBoundedEnum
         ((,) <$> genPermutation <#*#> Gen.matrix) applyToMatrix) :
   ("applyMultiply",
      checkForAllExtra QC.arbitraryBoundedEnum
         ((,,) <$> genPermutation <#*#> genPermutation <#*#> Gen.matrix)
         applyMultiply) :
   ("applyTranspose",
      checkForAllExtra QC.arbitraryBoundedEnum
         ((,) <$> genPermutation <#*#> Gen.matrix) applyTranspose) :
   ("determinantNumber",
      checkForAll genPermMatrix determinantNumber) :
   []