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

-- | generate a Sequence of non-nested intervals by means of 'fromEndPoints'
genNonNestedIntervalSeq :: Gen (Seq Intv)
genNonNestedIntervalSeq =
    withShrinkSeq $
    filterM (const QC.arbitrary) . fromEndPoints . List.sort
        =<< QC.listOf genUTCTime
-- TODO: these are also non-properly-overlapping, but we wish to include 
-- non-containment overlaps in the tests.

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