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 (ShapeInt, shapeInt, (#*##))
import Numeric.LAPACK.Vector (Vector)

import qualified Numeric.Netlib.Class as Class

import qualified Data.Array.Comfort.Shape as Shape

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

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


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.toSquare p)



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) :
   []