{-# LANGUAGE OverlappingInstances       #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE IncoherentInstances #-}

module Properties.Utils (
  module Properties.Utils,
  module Test.QuickCheck,
  module Test.QuickCheck.Batch,
  ) where

import Test.QuickCheck
import Test.QuickCheck.Batch
import Text.Show.Functions
import Control.Monad.Instances

import Control.Monad (liftM,liftM2,liftM5)

import qualified Data.Array.Vector as S
import Data.Array.Vector ((:*:)(..))

import Data.Word
import Data.Int
import Data.Complex
import Data.Ratio
import Data.List

opts = TestOptions {
         no_of_tests     = 500,
         length_of_tests = 0,
         debug_tests = False
       }

eq0 f g = property $
    model f                   == g
eq1 f g = \x               -> property $
    model (f x)               == g (model x)
eq2 f g = \x y             -> property $
    model (f x y)             == g (model x) (model y)
eq3 f g = \x y z           -> property $
    model (f x y z)           == g (model x) (model y) (model z)
eq4 f g = \x y z a         -> property $
    model (f x y z a)         == g (model x) (model y) (model z) (model a)
eq5 f g = \x y z a b       -> property $
    model (f x y z a b)       == g (model x) (model y) (model z) (model a) (model b)
eq6 f g = \x y z a b c     -> property $
    model (f x y z a b c)     == g (model x) (model y) (model z) (model a) (model b) (model c)
eq7 f g = \x y z a b c d   -> property $
    model (f x y z a b c d)   == g (model x) (model y) (model z) (model a) (model b) (model c) (model d)
eq8 f g = \x y z a b c d e -> property $
    model (f x y z a b c d e) == g (model x) (model y) (model z) (model a) (model b) (model c) (model d) (model e)

eqnotnull1 f g = \x     -> (not (S.nullU x)) ==> eq1 f g x
eqnotnull2 f g = \x y   -> (not (S.nullU y)) ==> eq2 f g x y
eqnotnull3 f g = \x y z -> (not (S.nullU z)) ==> eq3 f g x y z

{-
eqfinite1 f g = \x     -> forAll arbitrary $ \n -> Prelude.take n (f x)     == Prelude.take n (g x)
eqfinite2 f g = \x y   -> forAll arbitrary $ \n -> Prelude.take n (f x y)   == Prelude.take n (g x y)
eqfinite3 f g = \x y z -> forAll arbitrary $ \n -> Prelude.take n (f x y z) == Prelude.take n (g x y z)
-}

newtype A = A Int deriving (Eq, Show, Read, Arbitrary, S.UA)
newtype B = B Int deriving (Eq, Show, Read, Arbitrary, S.UA)
newtype C = C Int deriving (Eq, Show, Read, Arbitrary, S.UA)
type D = A
type E = B
type F = C
type G = A
type H = B

{-}
instance NatTrans S.UArr [] where
    eta = S.fromU
-}  
instance NatTrans S.MaybeS Maybe where
    eta (S.JustS a) = Just a
    eta S.NothingS = Nothing 

newtype OrdA = OrdA Int deriving (Eq, Ord, Show, Arbitrary, S.UA)

newtype N = N Int deriving (Eq, Ord, Num, Show, Arbitrary, S.UA)
newtype I = I Int deriving (Eq, Ord, Num, Enum, Real, Integral, Show, Arbitrary, S.UA)

instance Arbitrary Word where
    arbitrary = fmap fromIntegral (arbitrary :: Gen Int)
    coarbitrary = undefined
    
instance Arbitrary Word8 where
    arbitrary = fmap fromIntegral (arbitrary :: Gen Int)
    coarbitrary = undefined
    
instance Arbitrary Word16 where
    arbitrary = fmap fromIntegral (arbitrary :: Gen Int)
    coarbitrary = undefined
        
instance Arbitrary Word32 where
    arbitrary = fmap fromIntegral (arbitrary :: Gen Int)
    coarbitrary = undefined

instance Arbitrary Word64 where
    arbitrary = fmap fromIntegral (arbitrary :: Gen Integer)
    coarbitrary = undefined

instance Arbitrary Int8 where
    arbitrary = fmap fromIntegral (arbitrary :: Gen Int)
    coarbitrary = undefined

instance Arbitrary Int16 where
    arbitrary = fmap fromIntegral (arbitrary :: Gen Int)
    coarbitrary = undefined

instance Arbitrary Int32 where
    arbitrary = fmap fromIntegral (arbitrary :: Gen Int)
    coarbitrary = undefined

instance Arbitrary Int64 where
    arbitrary = fmap fromIntegral (arbitrary :: Gen Integer)
    coarbitrary = undefined                

instance (Arbitrary a, RealFloat a) => Arbitrary (Complex a) where
    arbitrary = liftM2 (:+) arbitrary arbitrary
    coarbitrary = undefined
    
instance (Arbitrary a, Integral a) => Arbitrary (Ratio a) where
    arbitrary = liftM2 (\x y -> x % if y == 0 then 1 else y) arbitrary arbitrary
    coarbitrary = undefined

instance Arbitrary Char where
    arbitrary     = elements ([' ', '\n', '\0'] ++ ['a'..'h'])
    coarbitrary c = variant (fromEnum c `rem` 4)

instance Arbitrary Ordering where
    arbitrary      = elements [LT, EQ, GT]
    coarbitrary LT = variant 0
    coarbitrary EQ = variant 1
    coarbitrary GT = variant 2

instance Arbitrary a => Arbitrary (S.MaybeS a) where
    arbitrary            = frequency [ (1, return S.NothingS)
                                     , (3, liftM S.JustS arbitrary) ]
    coarbitrary S.NothingS  = variant 0
    coarbitrary (S.JustS a) = variant 1 . coarbitrary a
    
instance (Arbitrary a, Arbitrary b) => Arbitrary (a :*: b) where
    arbitrary = do x <- arbitrary
                   y <- arbitrary
                   return ( x :*: y )
    coarbitrary (a:*:b) = coarbitrary a . coarbitrary b

instance (Arbitrary a, Arbitrary b, Arbitrary c, Arbitrary d, Arbitrary e)
      => Arbitrary (a, b, c, d ,e )
 where
  arbitrary = liftM5 (,,,,) arbitrary arbitrary arbitrary arbitrary arbitrary
  coarbitrary (a, b, c, d, e) =
    coarbitrary a . coarbitrary b . coarbitrary c . coarbitrary d .  coarbitrary e

instance (Arbitrary a, Arbitrary b, Arbitrary c, Arbitrary d, Arbitrary e, Arbitrary f)
      => Arbitrary (a, b, c, d, e, f)
 where
  arbitrary = liftM6 (,,,,,) arbitrary arbitrary arbitrary arbitrary arbitrary arbitrary
  coarbitrary (a, b, c, d, e, f) =
    coarbitrary a . coarbitrary b . coarbitrary c . coarbitrary d .  coarbitrary e . coarbitrary f

liftM6  :: (Monad m) => (a1 -> a2 -> a3 -> a4 -> a5 -> a6 -> r) -> m a1 -> m a2 -> m a3 -> m a4 -> m a5 -> m a6 -> m r
liftM6 f m1 m2 m3 m4 m5 m6 = do { x1 <- m1; x2 <- m2; x3 <- m3; x4 <- m4; x5 <- m5; x6 <- m6; return (f x1 x2 x3 x4 x5 x6) }

instance (Arbitrary a, Arbitrary b, Arbitrary c, Arbitrary d, Arbitrary e, Arbitrary f, Arbitrary g)
      => Arbitrary (a, b, c, d, e, f, g)
 where
  arbitrary = liftM7 (,,,,,,) arbitrary arbitrary arbitrary arbitrary arbitrary arbitrary arbitrary
  coarbitrary (a, b, c, d, e, f, g) =
    coarbitrary a . coarbitrary b . coarbitrary c . coarbitrary d .  coarbitrary e . coarbitrary f . coarbitrary g

liftM7  :: (Monad m) => (a1 -> a2 -> a3 -> a4 -> a5 -> a6 -> a7 -> r) -> m a1 -> m a2 -> m a3 -> m a4 -> m a5 -> m a6 -> m a7 -> m r
liftM7 f m1 m2 m3 m4 m5 m6 m7 = do { x1 <- m1; x2 <- m2; x3 <- m3; x4 <- m4; x5 <- m5; x6 <- m6; x7 <- m7 ; return (f x1 x2 x3 x4 x5 x6 x7) }


------------------------------------------------------------------------
-- Arbitrary instance for Stream

instance (S.UA a, Arbitrary a) => Arbitrary (S.UArr a) where
    arbitrary = do xs <- arbitrary
                   return $ S.toU xs
    coarbitrary = undefined

-- To let us generate two UArrs of equal length
data ELUArrs a b = ELUArrs !(S.UArr a) !(S.UArr b)
  deriving (Show, Eq)
  
instance (S.UA a, S.UA b, Arbitrary a, Arbitrary b) => Arbitrary (ELUArrs a b) where
    arbitrary = do n <- arbitrary
                   xs <- mapM (const arbitrary) $ replicate n 0
                   ys <- mapM (const arbitrary) $ replicate n 0
                   return $ ELUArrs (S.toU xs) (S.toU ys)
    coarbitrary = undefined
    
data ELUArrs3 a b c = ELUArrs3 !(S.UArr a) !(S.UArr b) !(S.UArr c)
  deriving (Show, Eq)

instance (S.UA a, S.UA b, S.UA c, Arbitrary a, Arbitrary b, Arbitrary c) => 
  Arbitrary (ELUArrs3 a b c) where
    arbitrary = do n <- arbitrary
                   xs <- mapM (const arbitrary) $ replicate n 0
                   ys <- mapM (const arbitrary) $ replicate n 0
                   zs <- mapM (const arbitrary) $ replicate n 0
                   return $ ELUArrs3 (S.toU xs) (S.toU ys) (S.toU zs)
    coarbitrary = undefined

data PosUArr = PosUArr !(S.UArr Int)
  deriving (Show, Eq)

instance Arbitrary PosUArr where
    arbitrary = do xs <- arbitrary
                   -- this isn't really uniform, but whatever
                   return $ PosUArr (S.toU . map abs $ xs)
    coarbitrary = undefined
    
data Ind2LenUArr a = Ind2LenUArr !(S.UArr a) !Int !Int !Int
  deriving (Show, Eq)
  
instance (Arbitrary a, S.UA a) => Arbitrary (Ind2LenUArr a) where
  arbitrary = do xs <- arbitrary
                 index1 <- fmap (`mod` (length xs + 1)) arbitrary -- TODO: check that this length + 1 stuff is correct
                 index2 <- fmap (`mod` (length xs + 1)) arbitrary
                 len <- fmap (`mod` (length xs - (max index1 index2) + 1)) arbitrary
                 return $ Ind2LenUArr (S.toU xs) index1 index2 len
  coarbitrary = undefined

data BoundedIndex a = BoundedIndex !(S.UArr a) !Int
    deriving (Show, Eq)
  
instance (Arbitrary a, S.UA a) => Arbitrary (BoundedIndex a) where
    arbitrary = do xs <- arbitrary
                   index <- fmap (`mod` (length xs + 1)) arbitrary
                   return $ BoundedIndex (S.toU xs) index
    coarbitrary = undefined

data CombineGen a = CombineGen !(S.UArr Bool) !(S.UArr a) !(S.UArr a) 
    deriving (Show, Eq)
    
instance (Arbitrary a, S.UA a) => Arbitrary (CombineGen a) where
    arbitrary = do fs <- arbitrary
                   -- don't want to depend on arrow, but I'm not sure why not
                   let (xl, yl) = (\(x, y) -> (length x, length y)) $ partition id fs
                   -- really ugly way to generate arbitraries of specific length
                   xs <- mapM (const arbitrary) [1..xl]
                   ys <- mapM (const arbitrary) [1..yl]
                   return $ CombineGen (S.toU fs) (S.toU xs) (S.toU ys)
    coarbitrary = undefined
    
{-
instance (Arbitrary a, Arbitrary s) => Arbitrary (S.Step a s)  where
    arbitrary = do x <- arbitrary
                   a <- arbitrary
                   s <- arbitrary
                   return $ case x of
                        LT -> S.Yield a s
                        EQ -> S.Skip s
                        GT -> S.Done
    coarbitrary = error "No coarbitrary for Step a s"
-}

-- existential state type
{-
instance (Arbitrary a) => Arbitrary (S.Stream a)  where
    coarbitrary = error "No coarbitrary for Streams"
    arbitrary = do xs    <- arbitrary :: Gen [a]
                   skips <- arbitrary :: Gen [Bool] -- random Skips
                   return (stream' (zip xs skips))
      where
        -- | Construct an abstract stream from a list, with Steps in it.
        stream' :: [(a,Bool)] -> S.Stream a
        stream' xs0 = S.Stream next (S.L xs0)
          where
            next (S.L [])             = S.Done
            next (S.L ((x,True ):xs)) = S.Yield x (S.L xs)
            next (S.L ((_,False):xs)) = S.Skip    (S.L xs)

instance Show a => Show (S.Stream a) where
  show = show . S.unstream

instance Eq a => Eq (S.Stream a) where
  xs == ys = S.unstream xs == S.unstream ys
-}

------------------------------------------------------------------------

class Model a b where
  model :: a -> b  -- get the abstract vale from a concrete value

instance S.UA a => Model (S.UArr a) [a] where model = S.fromU

instance S.UA a => Model (S.UArr a) (S.UArr a) where model = id
instance Model A A where model = id
instance Model B B where model = id
instance Model Bool Bool where model = id
instance Model Int  Int  where model = id
instance Model N    N    where model = id
instance Model OrdA OrdA where model = id
instance Model Ordering Ordering where model = id

instance (Model a a , Model b b) => Model (a:*:b) (a,b) where
        model (x:*:y) = (model x, model y)

-- not really moral
instance Functor ((:*:) a) where
        fmap f (x:*:y) = (x :*: f y)

-- More structured types are modeled recursively, using the NatTrans class from Gofer.
class (Functor f, Functor g) => NatTrans f g where
    eta :: f a -> g a

instance NatTrans [] []             where eta = id
instance NatTrans Maybe Maybe       where eta = id

instance NatTrans ((->) A) ((->) A) where eta = id
instance NatTrans ((->) B) ((->) B) where eta = id
instance NatTrans ((->) N) ((->) N) where eta = id
instance NatTrans ((->) C) ((->) C) where eta = id

instance Model f g => NatTrans ((,) f) ((,) g)
    where eta (f,a) = (model f, a)
instance Model f g => NatTrans ((:*:) f) ((:*:) g)
    where eta (f:*:a) = (model f:*: a)

instance (NatTrans m n, Model a b) => Model (m a) (n b)
    where model x = fmap model (eta x)