module Data.IntervalTest (
    genNonNestedIntervalSeq,
    genInterval,
    genIntervalSeq,
    genLabeledSeq,
    genNonEmptyInterval,
    genNonEmptyIntervalSeq,
    genSortedIntervals,
    genSortedIntervalSeq,
    genSortedList,
    forevery,
    foreveryPair,
    foreveryPairOf
    ) where

import Data.Interval

import qualified Data.Time as Time
import qualified Data.Sequence as Seq
import qualified Data.List as List
import Data.Foldable (toList)
import Data.Sequence (Seq,(|>))
import Data.Time (UTCTime)

import Control.Arrow (first)
import Control.Applicative (liftA2)

import qualified Test.QuickCheck as QC


type Intv = (UTCTime,UTCTime)

forevery :: (Show a, QC.Testable test) => Gen a -> (a -> test) -> QC.Property
forevery = uncurry QC.forAllShrink

foreveryPair ::
    (Show a, QC.Testable test) =>
    Gen a -> (a -> a -> test) -> QC.Property
foreveryPair (gen,shrink) =
    QC.forAllShrink
       (liftA2 (,) gen gen)
       (\(a,b) -> map (flip (,) b) (shrink a) ++ map ((,) a) (shrink b)) .
    uncurry

foreveryPairOf :: (Show a, Show b, QC.Testable test) =>
    Gen a -> Gen b -> (a -> b -> test) -> QC.Property
foreveryPairOf (genA,shrinkA) (genB,shrinkB) = 
    QC.forAllShrink
       (liftA2 (,) genA genB)
       (\(a,b) -> map (flip (,) b) (shrinkA a) ++ map ((,) a) (shrinkB b)) .
    uncurry

type Gen a = (QC.Gen a, a -> [a])

withoutShrink :: QC.Gen a -> Gen a
withoutShrink gen = (gen, const [])

shrinkList :: [a] -> [[a]]
shrinkList xs = List.zipWith (++) (List.inits xs) (List.tail $ List.tails xs)

shrinkSeq :: Seq a -> [Seq a]
shrinkSeq = map Seq.fromList . shrinkList . toList

withShrinkList :: QC.Gen [a] -> Gen [a]
withShrinkList gen = (gen, shrinkList)

withShrinkSeq :: QC.Gen (Seq a) -> Gen (Seq a)
withShrinkSeq gen = (gen, shrinkSeq)


genInterval :: Gen Intv
genInterval = withoutShrink $ do
    a <- genUTCTime
    b <- genUTCTime
    return (min a b, max a b)

genNonEmptyInterval :: Gen Intv
genNonEmptyInterval = withoutShrink $ do
    a <- genUTCTime
    b <- genUTCTime `QC.suchThat` (a/=)
    return (min a b, max a b)

genUTCTime :: QC.Gen UTCTime
genUTCTime = do
    day <- QC.arbitrary
    return $ Time.UTCTime (Time.ModifiedJulianDay day) 0

genIntervalSeq :: Gen (Seq Intv)
genIntervalSeq =
    withShrinkSeq $ fmap Seq.fromList $ QC.listOf $ fst genInterval

-- | For @i :: Intv@ generates @j@ such that 
-- @lb i < lb j@ and @ub i < ub j@ 
genNextIntv :: Intv -> QC.Gen Intv
genNextIntv i = do
    x <- genUTCTime `QC.suchThat` ((lb i)<)
    y <- genUTCTime `QC.suchThat` ((max x (ub i))<)
    return (x,y)

-- | generate a Sequence of non-nested intervals
genNonNestedIntervalSeq :: Gen (Seq Intv)
genNonNestedIntervalSeq = withShrinkSeq $ fst genInterval >>= go mempty where
    go js j = do
        done <- QC.arbitrary
        if done then return js else do
            i <- genNextIntv j
            go (js |> j) i

genNonEmptyIntervalSeq :: Gen (Seq Intv)
genNonEmptyIntervalSeq =
    withShrinkSeq $ fmap Seq.fromList $ QC.listOf $ fst genNonEmptyInterval

_genInterval :: Gen (Int,Int)
_genInterval = withoutShrink $ do
    a <- QC.arbitrary
    b <- QC.arbitrary
    return (min a b, max a b)

genSortedIntervals :: Gen [Intv]
genSortedIntervals =
    withShrinkList $
    fmap (List.sortBy (\i j -> compare (ub i) (ub j) <> compare (lb j) (lb i))) $ QC.listOf $ fst genInterval

genSortedIntervalSeq :: Gen (Seq Intv)
genSortedIntervalSeq = first (fmap sortByRight) genIntervalSeq

genSortedList :: Gen [Int]
genSortedList = withShrinkList $ fmap List.sort QC.arbitrary


genLabeledSeq :: Gen (Seq (Char,Intv))
genLabeledSeq =
    withShrinkSeq $
    mapM (liftA2 (,) (QC.choose ('a','z')) . pure) =<< fst genIntervalSeq