module Test.Array where

import qualified Data.Array.Knead.Parameterized.Render as Render
import qualified Data.Array.Knead.Simple.Symbolic as Symb
import qualified Data.Array.Knead.Simple.Slice as Slice
import qualified Data.Array.Knead.Expression as Expr
import qualified Data.Array.Knead.Shape as Shape
import qualified Data.Array.Comfort.Storable as Array
import qualified Data.Array.Comfort.Shape as ComfortShape
import Data.Array.Comfort.Storable (Array)

import qualified LLVM.Extra.Storable as Storable
import qualified LLVM.Extra.Marshal as Marshal
import qualified LLVM.Extra.Multi.Value as MultiValue

import qualified LLVM.Core as LLVM

import qualified Type.Data.Num.Decimal as TypeNum

import Foreign.Storable (Storable)

import qualified Data.List.HT as ListHT
import Data.Int (Int32, Int64)

import Control.Applicative ((<$>))

import qualified Test.QuickCheck.Monadic as QCMon
import qualified Test.QuickCheck as QC


type Dim = ComfortShape.ZeroBased Int64
type Dim2 = (Dim, Dim)

genArray :: (QC.Arbitrary a, Storable a) => QC.Gen (Array Dim2 a)
genArray = do
   m <- QC.choose (1,10)
   n <- QC.choose (1,10)
   let shape = (Shape.ZeroBased m, Shape.ZeroBased n)
   Array.fromList shape <$> QC.vector (ComfortShape.size shape)


rowSumSymb ::
   (Shape.C sh0, Shape.C sh1, MultiValue.Additive a) =>
   Symb.Array (sh0,sh1) a -> Symb.Array sh0 a
rowSumSymb = Symb.fold1 Expr.add

columnSumSymb ::
   (Shape.C sh0, Shape.C sh1, MultiValue.Additive a) =>
   Symb.Array (sh0,sh1) a -> Symb.Array sh1 a
columnSumSymb = Symb.fold1 Expr.add . Slice.apply Slice.transpose


getRows ::
   (ComfortShape.C sh0, ComfortShape.C sh1, Storable a) =>
   Array (sh0,sh1) a -> [[a]]
getRows x =
   ListHT.sliceVertical
      (ComfortShape.size $ snd $ Array.shape x)
      (Array.toList x)

rowPred ::
   (Num a, Eq a, Storable a,
    ComfortShape.C sh0, ComfortShape.C sh1) =>
   Array (sh0, sh1) a -> Array sh0 a -> Bool
rowPred x y  =  Array.toList y == map sum (getRows x)

columnPred ::
   (Num a, Eq a, Storable a,
    ComfortShape.C sh0, ComfortShape.C sh1) =>
   Array (sh0, sh1) a -> Array sh1 a -> Bool
columnPred x y  =  Array.toList y == foldl1 (zipWith (+)) (getRows x)

run ::
   (Shape.C sh0, Marshal.MV sh0, Show sh0,
    Shape.C sh1, Marshal.MV sh1, Show sh1,
    Show a, Num a, Eq a, Storable.C a) =>
   QC.Gen (Array sh0 a) ->
   (Symb.Array sh0 a -> Symb.Array sh1 a) ->
   (Array sh0 a -> Array sh1 a -> Bool) ->
   IO QC.Property
run qcgen code predicate = do
   act <- Render.run code
   return $ QC.forAll qcgen $ \x ->
      QCMon.monadicIO $ do
         y <- QCMon.run $ act x
         QCMon.assert $ predicate x y


tests :: [(String, IO QC.Property)]
tests =
   ("rowSum",
      run (genArray :: QC.Gen (Array Dim2 Int32)) rowSumSymb rowPred) :
   ("columnSum",
      run (genArray :: QC.Gen (Array Dim2 Int32)) columnSumSymb columnPred) :
   ("rowSumV3",
      run (genArray :: QC.Gen (Array Dim2 (LLVM.Vector TypeNum.D3 Int32)))
         rowSumSymb rowPred) :
   ("columnSumV3",
      run (genArray :: QC.Gen (Array Dim2 (LLVM.Vector TypeNum.D3 Int32)))
         columnSumSymb columnPred) :
   []