{-# OPTIONS_GHC -fglasgow-exts #-}
{-# LANGUAGE FlexibleContexts, FlexibleInstances, MultiParamTypeClasses, FunctionalDependencies, NoMonomorphismRestriction, UndecidableInstances #-}
import Control.Arrow
import Test.QuickCheck
import Text.Show.Functions
import Data.Array.CArray
import Data.Ix.Shapable (shapeToStride)
import Data.Array.Unboxed
import Data.Binary
import Data.List
import Foreign.Storable
import Text.Printf
import System.Environment (getArgs)
import System.IO
import System.Random

instance (Ix i, Arbitrary i, Storable e, Arbitrary e) => Arbitrary (CArray i e) where
    arbitrary = do
        a <- arbitrary
        b <- arbitrary
        let l = min a b
            u = max a b
        es <- vector (rangeSize (l,u))
        return $ listArray (l,u) es
    coarbitrary a = coarbitrary (assocs a)

instance (Ix i, Arbitrary i, Arbitrary e, IArray UArray e) => Arbitrary (UArray i e) where
    arbitrary = do
        a <- arbitrary
        b <- arbitrary
        let l = min a b
            u = max a b
        es <- vector (rangeSize (l,u))
        return $ listArray (l,u) es
    coarbitrary a = coarbitrary (assocs a)

class Model a b where model :: a -> b

instance (Ix i, IArray a e, Model i i', Model e e') => Model (a i e) ((i',i'),[e']) where
    model = (model . bounds &&& map model . elems)
instance (Model i i', Model e e', Ix i', IArray a e') => Model ((i,i),[e]) (a i' e') where
    model = uncurry listArray . (model *** map model)
instance (Ix i, Ix i', Model i i', Model e e', Storable e, IArray UArray e')
    => Model (CArray i e) (UArray i' e') where
    model = uncurry listArray . (model . bounds &&& map model . elems)
instance (Ix i, Ix i', Model i i', Model e e', Storable e', IArray UArray e)
    => Model (UArray i e) (CArray i' e') where
    model = uncurry listArray . (model . bounds &&& map model . elems)

-- Types are trivially modeled by themselves
instance Model Bool  Bool         where model = id
instance Model Int   Int          where model = id
instance Model Float Float        where model = id
instance Model Double Double      where model = id
instance (Model a a', Model b b') => Model (a,b) (a',b') where
    model (a,b) = (model a, model b)
instance (Model a a', Model b b', Model c c') => Model (a,b,c) (a',b',c') where
    model (a,b,c) = (model a, model b, model c)
instance (Model a a', Model b b', Model c c', Model d d') => Model (a,b,c,d) (a',b',c',d') where
    model (a,b,c,d) = (model a, model b, model c, model d)

f =|= g = \a         ->
    model (f a)         == g (model a)
f =||= g = \a b       ->
    model (f a b)       == g a (model b)
infix 1 =|=
infix 1 =||=

f =|||= g = \a b c     ->
    model (f a b c)     == g a (model b) c
eq4 f g = \a b c d   ->
    model (f a b c d)   == g (model a) (model b) (model c) (model d)
eq5 f g = \a b c d e ->
    model (f a b c d e) == g (model a) (model b) (model c) (model d) (model e)

(===) :: (Eq b) => (a -> b) -> (a -> b) -> a -> Bool
(f === g) x = f x == g x
infixl 1 ===

transposeArray a = ixmap ((swap *** swap) (bounds a)) swap a
    where swap = (\(i,j) -> (j,i))

prop_flatten_flatten = flatten . flatten === flatten
prop_reshape_flatten a = reshape (0, size a - 1) a == flatten a
prop_rank = length . shape === rank
prop_shape_size = product . shape === size
prop_size = size === rangeSize . bounds
prop_shape_stride_last = last . shapeToStride . shape === const 1
prop_transpose = transposeArray . transposeArray === id

ca_tests :: [(String, CArray (Int,Int) Double -> Bool)]
ca_tests = [ ("flatten flatten"   , prop_flatten_flatten)
           , ("reshape flatten"   , prop_reshape_flatten)
           , ("rank"              , prop_rank)
           , ("shape size"        , prop_shape_size)
           , ("size"              , prop_size)
           , ("shape stride last" , prop_shape_stride_last)
           , ("transpose^2"       , prop_transpose)
           ]

prop_amap =    (amap :: (Int -> Double) -> CArray Int Int -> CArray Int Double)
          =||= (amap :: (Int -> Double) -> UArray Int Int -> UArray Int Double)

prop_slice_all :: (Int -> Double) -> CArray (Int,Int) Int -> Property
prop_slice_all f a = size a > 0 ==> sliceWith (bounds a) (bounds a) f a == amap f a
prop_ixmapWithInd_amap :: (Int -> Double) -> CArray (Int,Int) Int -> Property
prop_ixmapWithInd_amap f a = size a > 0 ==> ixmapWithInd (bounds a) id (\_ e _ -> f e) a == amap f a

type Acc = Int
prop_accum f a ies = all (inRange (bounds a) . fst) ies
    ==> (      (accum :: (Int -> Acc -> Int) -> CArray Int Int -> [(Int, Acc)] -> CArray Int Int)
         =|||= (accum :: (Int -> Acc -> Int) -> UArray Int Int -> [(Int, Acc)] -> UArray Int Int)) f a ies

prop_composeAssoc f g h = (f . g) . h === f . (g . h)
    where types = [f,g,h] :: [CArray Int Int -> CArray Int Int]

main = do
    x <- getArgs
    let n = if null x then 100 else read . head $ x
        conf = Config { configMaxTest = n
                      , configMaxFail = 1000
                      , configSize = (+ 3) . (`div` 2)
                      , configEvery = \n args -> let s = show n in s ++ [ '\b' | _ <- s]
                      }
        mycheck (s,a) = printf "%-25s: " s >> check conf a
    mapM_ mycheck ca_tests
    mapM_ mycheck [ ("amap"        , prop_amap) ]
    mapM_ mycheck [ ("accum"       , prop_accum) ]
    mapM_ mycheck [ ("composeAssoc", prop_composeAssoc) ]
    mapM_ mycheck [ ("slice all"         , prop_slice_all)
                  , ("ixmapWithInd amap" , prop_ixmapWithInd_amap) ]

-- arb n k = generate n (mkStdGen k) arbitrary