module Data.IntervalTest (
    genDisjointIntervalSeq,
    genInterval,
    genIntervalSeq,
    genLabeledSeq,
    genNonEmptyInterval,
    genNonEmptyIntervalSeq,
    genSortedIntervals,
    genSortedIntervalSeq,
    genSortedList,
    (/\),
    (/\*),
    ) 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.Function (on)
import Data.Time (UTCTime)

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

import qualified Test.QuickCheck as QC


type Intv = (UTCTime,UTCTime)

infixr 0 /\, /\*

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

(/\*) ::
    (Show a, QC.Testable test) =>
    Gen a -> (a -> a -> test) -> QC.Property
(/\*) (gen,shrink) =
    QC.forAllShrink
       (liftA2 (,) gen gen)
       (\(a,b) -> map (flip (,) b) (shrink a) ++ map ((,) a) (shrink 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

genDisjointIntervalSeq :: Gen (Seq Intv)
genDisjointIntervalSeq =
    withShrinkSeq $
    filterM (const QC.arbitrary) . fromEndPoints . List.sort
        =<< QC.listOf genUTCTime

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 (compare `on` fst)) $ QC.listOf $ fst genInterval

genSortedIntervalSeq :: Gen (Seq Intv)
genSortedIntervalSeq = first (fmap sortByLeft) 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