{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ExistentialQuantification #-}
module Test.Shape where

import Test.Utility (genOrder, prefix)

import qualified Data.Array.Comfort.Shape.Test as ShapeTest
import qualified Data.Array.Comfort.Shape as Shape

import qualified Numeric.LAPACK.Matrix.Layout as Layout
import qualified Numeric.LAPACK.Matrix.Extent as Extent
import qualified Numeric.LAPACK.Permutation as Perm
import qualified Numeric.LAPACK.Shape as ExtShape
import Numeric.LAPACK.Matrix (ShapeInt, shapeInt)

import qualified Type.Data.Num.Unary as Unary
import Type.Data.Num.Unary (unary)

import Control.Applicative ((<$>))

import qualified Test.QuickCheck as QC


genPermutation :: QC.Gen (Perm.Shape ShapeInt)
genPermutation = Perm.Shape . Shape.ZeroBased <$> QC.choose (0,10)

genIntIndexed :: QC.Gen (ExtShape.IntIndexed (Shape.UpperTriangular ShapeInt))
genIntIndexed =
   ExtShape.IntIndexed . Shape.Triangular Shape.Upper . shapeInt
   <$>
   QC.choose (0,10)


genGeneral :: QC.Gen (Layout.General ShapeInt ShapeInt)
genGeneral = do
   order <- genOrder
   m <- QC.choose (0,10)
   n <- QC.choose (0,10)
   return $ Layout.general order (shapeInt m) (shapeInt n)

genTall :: QC.Gen (Layout.Tall ShapeInt ShapeInt)
genTall = do
   order <- genOrder
   m <- QC.choose (0,10)
   n <- QC.choose (0,m)
   return $ Layout.tall order (shapeInt m) (shapeInt n)

genWide :: QC.Gen (Layout.Wide ShapeInt ShapeInt)
genWide = do
   order <- genOrder
   m <- QC.choose (0,10)
   n <- QC.choose (m,10)
   return $ Layout.wide order (shapeInt m) (shapeInt n)

genSquare :: QC.Gen (Layout.Square ShapeInt)
genSquare = do
   order <- genOrder
   n <- QC.choose (0,10)
   return $ Layout.square order (shapeInt n)


genDiagonal :: QC.Gen (Layout.Diagonal ShapeInt)
genDiagonal = do
   order <- genOrder
   n <- QC.choose (0,10)
   return $ Layout.diagonal order (shapeInt n)

genLowerTriangular ::
   Layout.PackingSingleton pack ->
   QC.Gen (Layout.LowerTriangularP pack ShapeInt)
genLowerTriangular pack = do
   order <- genOrder
   n <- QC.choose (0,10)
   return $ Layout.lowerTriangularP pack order (shapeInt n)

genUpperTriangular ::
   Layout.PackingSingleton pack ->
   QC.Gen (Layout.UpperTriangularP pack ShapeInt)
genUpperTriangular pack = do
   order <- genOrder
   n <- QC.choose (0,10)
   return $ Layout.upperTriangularP pack order (shapeInt n)

genSymmetric ::
   Layout.PackingSingleton pack ->
   QC.Gen (Layout.SymmetricP pack ShapeInt)
genSymmetric pack = do
   order <- genOrder
   n <- QC.choose (0,10)
   return $ Layout.symmetricP pack order (shapeInt n)

genHermitian ::
   Layout.PackingSingleton pack ->
   QC.Gen (Layout.HermitianP pack ShapeInt)
genHermitian pack = do
   order <- genOrder
   n <- QC.choose (0,10)
   return $ Layout.hermitianP pack order (shapeInt n)


data Banded meas vert horiz height width =
   forall sub super.
   (Unary.Natural sub, Unary.Natural super) =>
   Banded (Layout.Banded sub super meas vert horiz height width)

instance
   (Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Show height, Show width, Shape.C height, Shape.C width) =>
      Show (Banded meas vert horiz height width) where
   showsPrec p (Banded sh) = showsPrec p sh

instance
   (Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Shape.C height, Shape.C width) =>
      Shape.C (Banded meas vert horiz height width) where

   size (Banded sh) = Shape.size sh

instance
   (Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Shape.Indexed height, Shape.Indexed width) =>
      Shape.Indexed (Banded meas vert horiz height width) where
   type Index (Banded meas vert horiz height width) =
            Layout.BandedIndex (Shape.Index height) (Shape.Index width)
   indices (Banded sh) = Shape.indices sh
   unifiedOffset (Banded sh) = Shape.unifiedOffset sh
   inBounds (Banded sh) = Shape.inBounds sh

   unifiedSizeOffset (Banded sh) = Shape.unifiedSizeOffset sh

instance
   (Extent.Measure meas, Extent.C vert, Extent.C horiz,
    Shape.InvIndexed height, Shape.InvIndexed width) =>
      Shape.InvIndexed (Banded meas vert horiz height width) where

   unifiedIndexFromOffset (Banded sh) = Shape.unifiedIndexFromOffset sh


genBanded ::
   Layout.Full meas vert horiz height width ->
   QC.Gen (Banded meas vert horiz height width)
genBanded sh = do
   kl <- QC.choose (0,10)
   ku <- QC.choose (0,10)
   Unary.reifyNatural kl $ \sub ->
      Unary.reifyNatural ku $ \super ->
      return $ Banded $ Layout.bandedFromFull (unary sub, unary super) sh


data BandedHermitian size =
   forall offDiag.
   (Unary.Natural offDiag) =>
   BandedHermitian (Layout.BandedHermitian offDiag size)

instance (Show size, Shape.C size) => Show (BandedHermitian size) where
   showsPrec p (BandedHermitian sh) = showsPrec p sh

instance (Shape.C size) => Shape.C (BandedHermitian size) where
   size (BandedHermitian sh) = Shape.size sh

instance (Shape.Indexed size) => Shape.Indexed (BandedHermitian size) where
   type Index (BandedHermitian size) =
            Layout.BandedIndex (Shape.Index size) (Shape.Index size)
   indices (BandedHermitian sh) = Shape.indices sh
   unifiedOffset (BandedHermitian sh) = Shape.unifiedOffset sh
   inBounds (BandedHermitian sh) = Shape.inBounds sh

   unifiedSizeOffset (BandedHermitian sh) = Shape.unifiedSizeOffset sh

instance
   (Shape.InvIndexed size) => Shape.InvIndexed (BandedHermitian size) where

   unifiedIndexFromOffset (BandedHermitian sh) =
      Shape.unifiedIndexFromOffset sh


genBandedHermitian :: QC.Gen (BandedHermitian ShapeInt)
genBandedHermitian = do
   order <- genOrder
   n <- QC.choose (0,10)
   k <- QC.choose (0,10)
   Unary.reifyNatural k $ \numOff ->
      return $ BandedHermitian $
         Layout.bandedHermitian (unary numOff) order (shapeInt n)


mosaicTests :: Layout.PackingSingleton pack -> [(String, QC.Property)]
mosaicTests pack =
   prefix ("LowerTriangular." ++ show pack)
      (ShapeTest.tests $ genLowerTriangular pack) ++
   prefix ("UpperTriangular." ++ show pack)
      (ShapeTest.tests $ genUpperTriangular pack) ++
   prefix ("Symmetric." ++ show pack)
      (ShapeTest.tests $ genSymmetric pack) ++
   prefix ("Hermitian." ++ show pack)
      (ShapeTest.tests $ genHermitian pack) ++
   []

tests :: [(String, QC.Property)]
tests =
   prefix "Permutation" (ShapeTest.tests genPermutation) ++
   prefix "IntIndexed" (ShapeTest.tests genIntIndexed) ++

   prefix "General" (ShapeTest.tests genGeneral) ++
   prefix "Tall" (ShapeTest.tests genTall) ++
   prefix "Wide" (ShapeTest.tests genWide) ++
   prefix "Square" (ShapeTest.tests genSquare) ++

   prefix "Split.Reflector.General"
      (ShapeTest.tests $
       Layout.splitFromFull Layout.Reflector <$> genGeneral) ++
   prefix "Split.Reflector.Tall"
      (ShapeTest.tests $
       Layout.splitFromFull Layout.Reflector <$> genTall) ++
   prefix "Split.Reflector.Wide"
      (ShapeTest.tests $
       Layout.splitFromFull Layout.Reflector <$> genWide) ++
   prefix "Split.Reflector.Square"
      (ShapeTest.tests $
       Layout.splitFromFull Layout.Reflector <$> genSquare) ++
   prefix "Split.Triangle.General"
      (ShapeTest.tests $
       Layout.splitFromFull Layout.Triangle <$> genGeneral) ++
   prefix "Split.Triangle.Tall"
      (ShapeTest.tests $
       Layout.splitFromFull Layout.Triangle <$> genTall) ++
   prefix "Split.Triangle.Wide"
      (ShapeTest.tests $
       Layout.splitFromFull Layout.Triangle <$> genWide) ++
   prefix "Split.Triangle.Square"
      (ShapeTest.tests $
       Layout.splitFromFull Layout.Triangle <$> genSquare) ++

   mosaicTests Layout.Packed ++
   mosaicTests Layout.Unpacked ++

   prefix "Diagonal" (ShapeTest.tests genDiagonal) ++
   prefix "Banded.General" (ShapeTest.tests $ genBanded =<< genGeneral) ++
   prefix "Banded.Tall" (ShapeTest.tests $ genBanded =<< genTall) ++
   prefix "Banded.Wide" (ShapeTest.tests $ genBanded =<< genWide) ++
   prefix "Banded.Square" (ShapeTest.tests $ genBanded =<< genSquare) ++
   prefix "BandedHermitian" (ShapeTest.tests genBandedHermitian) ++
   []