module Test.Permutation where

import qualified Test.Generator as Gen
import qualified Test.Utility as Util
import Test.Generator ((<|*|>))
import Test.Utility (equalArray, Tagged(Tagged))

import qualified Numeric.LAPACK.Permutation as Perm
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 (ZeroInt, zeroInt, (<#>))
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 Foreign.C.Types (CInt)

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

import Data.Monoid ((<>))

import qualified Test.QuickCheck as QC


genPivots :: [()] -> QC.Gen (Vector ZeroInt CInt)
genPivots nat = do
   let n = length nat
   let nc = fromIntegral n
   fmap (Vector.fromList (zeroInt n)) $
      forM (zip [1..] nat) $ \(i,()) -> QC.choose (i,nc)

genPerm :: [()] -> QC.Gen (Permutation ZeroInt)
genPerm = fmap (\p -> Perm.fromPivots NonInverted (Array.shape p) p) . genPivots


permutationPivots :: Inversion -> Vector ZeroInt CInt -> Bool
permutationPivots inv xs =
   Array.toList (Perm.toPivots inv (Perm.fromPivots inv (Array.shape xs) xs))
   ==
   Array.toList xs

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

transposeMultiply :: (Permutation ZeroInt, Permutation ZeroInt) -> 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 :: Gen.Matrix a Int Int (Permutation ZeroInt)
genPermutation =
   flip Gen.mapGen Gen.squareDim $ \_ sh@(Shape.ZeroBased n) ->
   let nc = fromIntegral n
   in fmap (Perm.fromPivots NonInverted sh . Vector.fromList sh) $
         forM [1..] $ \i -> QC.choose (i,nc)


permApply ::
   (Shape.C height, Eq height, Shape.C width, Class.Floating a) =>
   Inversion -> Permutation height ->
   Matrix.General height width a ->
   Matrix.General height width a
permApply inv = Perm.apply (inv==Inverted)

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 ZeroInt, Matrix.General ZeroInt ZeroInt a) -> Bool
applyToMatrix inv (p,m) =
   equalArray
      (permApply inv p m)
      (permToMatrix inv p <#> m)

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

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


addTag :: Gen.T a dim array -> Gen.T a dim (Tagged a array)
addTag = fmap Tagged

taggedToMatrix ::
   (Class.Floating a) => Tagged a (Permutation ZeroInt) -> Square ZeroInt a
taggedToMatrix (Tagged p) = Perm.toMatrix p

determinantNumber ::
   (Class.Floating a, Eq a) => Tagged a (Permutation ZeroInt) -> Bool
determinantNumber tp@(Tagged p) =
   Perm.numberFromSign (Perm.determinant p)
   ==
   Square.determinant (taggedToMatrix tp)


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 tag dim 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 tag dim 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 =
   ("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 (addTag genPermutation) determinantNumber) :
   []