{-# OPTIONS_GHC -fno-warn-orphans -XFlexibleInstances #-}


module ArbTest (
   ArbEvents (..),
   RtspTest (..),
   interpret,
   compile,
   prop_emitsBefore,
   prop_RTSP,
   prop_rtspMonoid1,
   prop_rtspMonoid2,
   prop_rtspMonoidCommutes,
   prop_rtspMonoidAssociates,
   prop_rtspCategoryId1,
   prop_rtspCategoryId2,
   prop_rtspCategoryAssociates,
   prop_rtspCategoryAssociates2,
   prop_rtspIfThenElse,
   prop_eventCount,
   prop_eventLatch
) where


-- import Control.Applicative
import Control.RTSP
import Control.Category
import Control.Monad
import Data.Function (on)
import Data.List (groupBy, sortBy, partition)
import Data.Maybe
import Data.Monoid
import Data.Ord
import Data.Time
import Test.QuickCheck


import Prelude hiding ((.), id)



-- Orphan instances for RTSP types

instance Arbitrary a => Arbitrary (Event a) where
   arbitrary = do
         dt <- arbitrary
         v <- arbitrary
         return $ Event (dt `addUTCTime` epoch) v
   shrink = shrinkNothing


instance Arbitrary NominalDiffTime where
   arbitrary = fmap (fromRational . abs) arbitrary
   shrink dt = if dt == dtSecs then [] else [dtSecs]
      where dtSecs = fromInteger $ floor dt 



-- | Base time for events.
epoch :: UTCTime
epoch = UTCTime (fromGregorian 2000 1 1) 0


-- | A list of arbitrary events in chronological order
newtype ArbEvents a = ArbEvents [Event a]

instance (Show a) => Show (ArbEvents a) where
   show (ArbEvents evs) = "[\n" ++ concatMap showEvent evs ++ "]"
      where showEvent (Event t v) = "   Event " ++ show (diffUTCTime t epoch) ++ " " ++ show v ++ "\n"

instance (Arbitrary a) => Arbitrary (ArbEvents a) where
   arbitrary = do
      times <- fmap (scanl (flip addUTCTime) epoch) arbitrary
      events <- forM times $ \t -> fmap (Event t) arbitrary
      return $ ArbEvents events
   shrink (ArbEvents evs) = map ArbEvents $ shrink evs


-- | An RTSP test consists of a descriptive string, a list function and an RTSP.  The list function
-- has the same effect on the list of events that the RTSP has on a stream.
data RtspTest a =
      Id
      | Delay Rational
      | Dup Rational
      | Func (a -> a) String
      | If (a -> Bool) String (RtspTest a) (RtspTest a)
      | Pipe (RtspTest a) (RtspTest a)
      | Par (RtspTest a) (RtspTest a)

instance Show (RtspTest a) where
   show Id = "id"
   show (Delay t) = "delay " ++ show t
   show (Dup t) = "duplicate " ++ show t
   show (Func _ str) = "stream " ++ str
   show (If _ str r1 r2) = "(ifThenElse (" ++ str ++ ", " ++ show r1 ++ ") (" ++ show r2 ++ "))"
   show (Pipe r1 r2) = "(" ++ show r1 ++ " >>> " ++ show r2 ++ ")"
   show (Par r1 r2) = "(" ++ show r1 ++ " `mappend` " ++ show r2 ++ ")"


instance (Integral a) => Arbitrary (RtspTest a) where
   arbitrary = do
      frequency [
         (2,  return Id),
         (1,  fmap (Delay . fromRational . abs) arbitrary),
         (1,  fmap (Dup . fromRational . abs) arbitrary),
         (1,  oneof [return $ Func (*2) "(*2)",
                     return $ Func (+1) "(+1)",
                     return $ Func (*3) "(*3)"]),
         (1, do
                 (p, str) <- elements [(odd, "odd"),
                                       (even, "even"),
                                       ((== 0) . (`mod` 3), "mult3")]
                 return (If p str) `ap` arbitrary `ap` arbitrary),
         (1, return Pipe `ap` arbitrary `ap` arbitrary),
         (1, return Par `ap` arbitrary `ap` arbitrary)
         ] 
   shrink Id = []
   shrink (Delay dt) = Id : map Delay (shrink dt)
   shrink (Dup dt) = Delay dt : map Dup (shrink dt)
   shrink (Func _ _) = [Id]
   shrink (If p str r1 r2) = shrinkBinaryOp (If p str) r1 r2
   shrink (Pipe r1 r2) = shrinkBinaryOp Pipe r1 r2
   shrink (Par r1 r2) = shrinkBinaryOp Par r1 r2


shrinkBinaryOp :: (Integral a) => (RtspTest a -> RtspTest a -> RtspTest a) -> RtspTest a -> RtspTest a -> [RtspTest a] 
shrinkBinaryOp op r1 r2 = concat [
   [Id, r1, r2], 
   map (\r -> op r r2) $ shrink r1,
   map (op r1) $ shrink r2]


-- | Interpret a test on a list, predicting the output for the equivalent arrow.
interpret :: (Num a) => RtspTest a -> [Event a] -> [Event a]
interpret Id             evs = evs
interpret (Delay dt)     evs = map (\(Event t v) -> Event (addUTCTime (fromRational dt) t) v) evs
interpret (Dup dt)       evs =
   foldl merge [] $ map (\(Event t v) -> [Event t v, Event (addUTCTime (fromRational dt) t) v]) evs
interpret (Func f _)     evs = map (fmap f) evs
interpret (If p _ r1 r2) evs = merge (interpret r1 thens) (interpret r2 elses)
   where 
      (thens, elses) = partition (p . eventValue) evs
interpret (Pipe r1 r2) evs = interpret r2 $ interpret r1 evs
interpret (Par r1 r2) evs = merge (interpret r1 evs) (interpret r2 evs)



-- | Merge two sorted lists of events.
merge :: [Event a] -> [Event a] -> [Event a]
merge xs [] = xs
merge [] ys = ys
merge xs@(x:xs1) ys@(y:ys1) = if y `isBefore` x then y : merge ys1 xs else x : merge ys xs1



-- | Compile a test into an arrow.
compile :: (Num a) => RtspTest a -> RTSP a a
compile Id = id
compile (Delay dt) = delay (fromRational dt)
compile (Dup dt) = accumulate $ repeatEvent [0, fromRational dt]
compile (Func f _) = stream f
compile (If p _ r1 r2) = ifThenElse (p, compile r1) (compile r2)
compile (Pipe r1 r2) = compile r1 >>> compile r2
compile (Par r1 r2) = compile r1 `mappend` compile r2


-- | Two event streams are equivalent regardless of the ordering of simultaneous events
isEquivalent :: (Ord a) => [Event a] -> [Event a] -> Bool
isEquivalent xs ys = normalise xs == normalise ys
   where normalise = map (sortBy (comparing eventValue)) . groupBy ((==) `on` eventTime)


-- | Assert that "compile" and "interpret" are equivalent.
prop_RTSP :: RtspTest Integer -> ArbEvents Integer -> Property
prop_RTSP tst (ArbEvents evs) = printTestCase failStr $ result1 `isEquivalent` result2
   where
      result1 = interpret tst evs
      result2 = simulateRTSP (compile tst) evs
      failStr = concat [  "interpret => ", show (ArbEvents result1), "\n",
                          "compile => ", show (ArbEvents result2), "\n"]
                          



-- | Reify primitive RTA actions.
data RtaTest s c =
   Modify (s -> s) | Emit (s -> c) | Pause (s -> NominalDiffTime)


instance (Arbitrary s, CoArbitrary s, Arbitrary c) => Arbitrary (RtaTest s c)
   where
      arbitrary = frequency [
         (3, return Modify `ap` arbitrary),
         (1, return Emit `ap` arbitrary),
         (5, return Pause `ap` arbitrary)
         ]
      shrink = shrinkNothing

instance Show (RtaTest s c) where
   show (Modify _) = "Modify"
   show (Emit _) = "Emit"
   show (Pause _) = "Pause"



instance (Integral s, Arbitrary s, CoArbitrary s, Arbitrary c) => Arbitrary (RTA s c Bool) where
   arbitrary = fmap execRtaTests arbitrary
   shrink = shrinkNothing

instance (CoArbitrary b, Arbitrary c) => Arbitrary (RTSP b c) where
   arbitrary = do
      rtaF <- arbitrary
      return $ execRTA (0 :: Integer) rtaF
   shrink = shrinkNothing


-- | Execute an RtaTest      
execRtaTest :: RtaTest s c -> RTA s c ()
execRtaTest (Modify f) = fmap f get >>= put
execRtaTest (Emit f) = fmap f get >>= emit
execRtaTest (Pause f) = fmap f get >>= pause


-- | Execute a sequence of RtaTests as a single action.
execRtaTests :: (Integral s) => [RtaTest s c] -> RTA s c Bool
execRtaTests ts = do
   mapM_ execRtaTest ts
   s <- get
   -- return True
   return $ (s `mod` 20) /= 0


type RtspProp = ArbEvents Integer -> Property

type RII = RTSP Integer Integer

rtspEquivalent :: (CoArbitrary b, Arbitrary c, Ord c, Show c) => RTSP b c -> RTSP b c -> ArbEvents b -> Property
rtspEquivalent r1 r2 (ArbEvents evs) = printTestCase failStr $ trace r1 `isEquivalent` trace r2
   where
      trace r = simulateRTSP r evs
      failStr = concat [
         "Trace1 = ", show $ ArbEvents $ trace r1, "\n",
         "Trace2 = ", show $ ArbEvents $ trace r2, "\n"]
         
         
prop_emitsBefore :: RII -> RII -> Event Integer -> Property
prop_emitsBefore r1 r2 ev =
   printTestCase (show (evs1, evs2)) $
   case (evs1, evs2) of
       (Nothing, Nothing) -> not (es1 `emitsBefore` es2 || es2 `emitsBefore` es1)
       (Just _,  Nothing) -> es1 `emitsBefore` es2 && not (es2 `emitsBefore` es1)
       (Nothing, Just _ ) -> not (es1 `emitsBefore` es2) && es2 `emitsBefore` es1
       (Just e1, Just e2) -> (e1 `isBefore` e2) == (es1 `emitsBefore` es2) &&
                             (e2 `isBefore` e1) == (es2 `emitsBefore` es1)
   where
      es1 = runRTSP r1 ev
      es2 = runRTSP r2 ev
      evs1 = listToMaybe $ esPeek es1
      evs2 = listToMaybe $ esPeek es2

prop_rtspMonoid1 :: RII -> RtspProp 
prop_rtspMonoid1 r = rtspEquivalent r (mempty `mappend` r)

prop_rtspMonoid2 :: RII -> RtspProp
prop_rtspMonoid2 r = rtspEquivalent r (r `mappend` mempty)


prop_rtspMonoidCommutes :: RII -> RII -> RtspProp
prop_rtspMonoidCommutes r1 r2 = rtspEquivalent (r1 `mappend` r2) (r2 `mappend` r1)

prop_rtspMonoidAssociates :: RII -> RII -> RII -> RtspProp
prop_rtspMonoidAssociates r1 r2 r3 = rtspEquivalent (r1 `mappend` (r2 `mappend` r3)) ((r1 `mappend` r2) `mappend` r3)


prop_rtspCategoryId1 :: RII -> RtspProp
prop_rtspCategoryId1 r = rtspEquivalent r (id >>> r)

prop_rtspCategoryId2 :: RII -> RtspProp
prop_rtspCategoryId2 r = rtspEquivalent r (r >>> id)


prop_rtspCategoryAssociates :: RII -> RII -> RII -> RtspProp
prop_rtspCategoryAssociates r1 r2 r3 = rtspEquivalent (r1 >>> (r2 >>> r3)) ((r1 >>> r2) >>> r3)

prop_rtspCategoryAssociates2 :: RtspTest Integer -> RtspTest Integer -> RtspTest Integer -> RtspProp
prop_rtspCategoryAssociates2 rt1 rt2 rt3 = rtspEquivalent (r1 >>> (r2 >>> r3)) ((r1 >>> r2) >>> r3)
   where
      r1 = compile rt1
      r2 = compile rt2
      r3 = compile rt3


-- | @ifThenElse (p, rThen) rElse@ is equivalent to
--
-- >  streamFilter (p, rThen) `mappend` streamFilter (not . p, rElse)
prop_rtspIfThenElse :: RII -> RII -> RtspProp
prop_rtspIfThenElse r1 r2 = rtspEquivalent
   (ifThenElse (odd, r1) r2)
   (streamFilter (odd, r1) `mappend` streamFilter (not . odd, r2))
--    where
--       r1 = compile rt1
--       r2 = compile rt2
   

traceEquivalent :: (Ord a, Show a) => [Event a] -> [Event a] -> Property
traceEquivalent trace1 trace2 = 
   printTestCase failStr $ trace1 `isEquivalent` trace2
   where failStr = concat [
            "Trace 1 = ", show $ ArbEvents trace1, "\n",
            "Trace 2 = ", show $ ArbEvents trace2, "\n"]
      
      
-- | Count events.      
eventCount :: RTSP b (Integer, b)
eventCount = execRTA 0 $ \v -> do
   s <- fmap (+1) get
   put s
   emit (s, v)
   return True

prop_eventCount :: RtspProp
prop_eventCount (ArbEvents evs) = traceEquivalent
   (simulateRTSP eventCount evs) 
   (zipWith (\n -> fmap (\v -> (n, v))) [1..] evs)



-- | Repeat each event value once a second ten times.
eventLatch :: RTSP b b
eventLatch = accumulateRTA () $ \v -> do
   replicateM_ 5 (emit v >> pause 1)
   return True
   
prop_eventLatch :: RtspProp
prop_eventLatch (ArbEvents evs) = traceEquivalent
   (simulateRTSP eventLatch evs)
   (sortBy (comparing eventTime) $ concatMap rep evs)
   where
      rep (Event t v) = [Event (n `addUTCTime` t) v | n <- [0,1..4]]