module Main where

-- Prelude

import Prelude hiding (id, (.))

-- Control

import Control.Applicative

-- Data

import           Data.Foldable (toList)
import           Data.MultiChange (MultiChange)
import qualified Data.MultiChange               as MultiChange
import           Data.Map (Map)
import qualified Data.Map                       as Map
import           Data.Incremental               
import qualified Data.Incremental.Tuple         as Tuple
import qualified Data.Incremental.Map           as IncMap

-- Criterion
import Criterion.Main
import Criterion.Types
import Control.DeepSeq
import GHC.Generics hiding (C)


-- Test

import Test.QuickCheck
import Test.QuickCheck.Poly

instance Arbitrary a => Arbitrary (PrimitiveChange a) where

    arbitrary = frequency [(1, keepGen), (5, replaceGen)] where

        keepGen = return Keep

        replaceGen = fmap ReplaceBy arbitrary

    shrink Keep            = []
    shrink (ReplaceBy val) = Keep : map ReplaceBy (shrink val)


instance Arbitrary p => Arbitrary (MultiChange p) where

    arbitrary = fmap MultiChange.singleton arbitrary

    shrink change = map MultiChange.fromList (shrink (toList change))


-- Pair changes
deriving instance (Show (DefaultChange a), Show (DefaultChange b)) =>
                  Show (Tuple.AtomicChange a b)

instance (Arbitrary (DefaultChange a), Arbitrary (DefaultChange b)) =>
         Arbitrary (Tuple.AtomicChange a b) where

    arbitrary = oneof [firstGen, secondGen] where

        firstGen = fmap Tuple.First arbitrary

        secondGen = fmap Tuple.Second arbitrary

    shrink (Tuple.First change)  = map Tuple.First (shrink change)
    shrink (Tuple.Second change) = map Tuple.Second (shrink change)

-- Map changes
instance (Arbitrary k, Arbitrary (DefaultChange k),Arbitrary v, Arbitrary (DefaultChange v)) =>
         Arbitrary (IncMap.AtomicChange k v) where

    arbitrary = oneof [delete, insert] where

        delete = liftA IncMap.Delete arbitrary

        insert = liftA2 IncMap.Insert arbitrary arbitrary

    shrink (IncMap.Insert k v)
        = [IncMap.Insert k' v'
              | (k', v') <- shrink (k, v)]
    shrink (IncMap.Delete k)
        = [IncMap.Delete k'
              | k' <- shrink k]

newtype AtomicAChange = DoubleAndAdd Integer deriving (Show, Arbitrary, Generic)

instance Change AtomicAChange where

    type Value AtomicAChange = A

    DoubleAndAdd diff $$ A integer = A (2 * integer + diff)

instance Changeable A where

    type DefaultChange A = MultiChange AtomicAChange

deriving instance Ord A

newtype AtomicBChange = TripleAndAdd Integer deriving (Show, Arbitrary, Generic)

instance Change AtomicBChange where

    type Value AtomicBChange = B

    TripleAndAdd diff $$ B integer = B (3 * integer + diff)

instance Changeable B where

    type DefaultChange B = MultiChange AtomicBChange

deriving instance Ord B

newtype AtomicCChange = QuadrupleAndAdd Integer deriving (Show, Arbitrary, Generic)

instance Change AtomicCChange where

    type Value AtomicCChange = C

    QuadrupleAndAdd diff $$ C integer = C (4 * integer + diff)

instance Changeable C where

    type DefaultChange C = MultiChange AtomicCChange

deriving instance Ord C


instance (NFData a) => NFData (MultiChange a)
instance NFData (Tuple.AtomicChange (Map A B) (Map A B))
instance NFData (Tuple.AtomicChange A B)
instance (Generic a, Generic b, NFData a, NFData b) => NFData (IncMap.AtomicChange a b)
instance NFData AtomicBChange
instance NFData AtomicAChange
instance NFData A
instance NFData B
instance NFData C

deriving instance (Generic a, Generic b) => Generic (IncMap.AtomicChange a b)
deriving instance Generic A
deriving instance Generic B
deriving instance Generic C

applyChanges :: Change p => (Value p, [p]) -> [Value p]
applyChanges (val, changes) = scanl (flip ($$)) val changes


testMap r n = do
    pairs <- generate $ vectorOf n $ resize r $ arbitrary 
    let resultMap = Map.fromList pairs
    return resultMap

--testMapSize :: Int -> IO (Int,Int)
--testMapSize n = do
--    pairs <- generate $ vectorOf n $ resize n $ (arbitrary :: Gen (A,B)) 
--    let resultMap = Map.fromList pairs
--    let resultMapSize = Map.size resultMap
--    return (n,resultMapSize)

testChanges r n = do
    changes <- generate $ vectorOf n (resize r $ arbitrary) 
    return changes

testFun r = do
    fun <- generate $ resize r $ arbitrary
    return fun

testKey r = do
    key <- generate $ resize r $ arbitrary
    return key

checkTrans trans (val,changes) = applyChanges (runTrans trans (val,changes))
checkOriginal trans (val,changes) = map trans (applyChanges (val,changes))

--benchmarks
main :: IO ()
main = 
  defaultMainWith
    (defaultConfig {reportFile = Just "performance.html"}) $ concat [
    map (benchFilter 10) [1000,10000,100000,1000000],
    map (benchMap 10) [1000,10000,100000,1000000],
    map (benchLookup 10) [1000,10000,100000,1000000],
    map (benchMember 10) [1000,10000,100000,1000000],
    map (benchPartition 10) [1000,10000,100000,1000000],
    map (benchUnion 10) [1000,10000,100000,1000000],
    map (benchDifference 10) [1000,10000,100000,1000000],
    map (benchIntersection 10) [1000,10000,100000,1000000],
    map (benchKeysSet 10) [1000,10000,100000,1000000],
    map (benchSubmapOf 10) [1000,10000,100000,1000000],
    map (benchSubmapOf 10) [1000,10000,100000,1000000],
    map (benchSplit 10) [1000,10000,100000,1000000],
    map (benchMapKeys 10) [1000,10000,100000,1000000]
    ]    



benchFilter ncgh n = env (do
            tMap <- testMap n n :: IO (Map A B) 
            tChanges <- testChanges n ncgh :: IO [MultiChange (IncMap.AtomicChange A B)]
            let tChangesScan = scanl (\x y -> x ++ [y]) [] tChanges
            let testValAndChanges = zip (replicate (ncgh+1) tMap) tChangesScan
            tFun  <- testFun n :: IO (B -> Bool)
            return (testValAndChanges,tFun))
         (\ ~(valAndChanges,fun) ->
            bgroup ("filter " ++ show n) $ concat 
            [map (\x -> bench ("original change=" ++ show x) $ nf (checkOriginal (Map.filter fun)) (valAndChanges !! x)) [0..ncgh],
             map (\x -> bench ("incremental change=" ++ show x) $ nf (checkTrans (IncMap.filter fun)) (valAndChanges !! x)) [0..ncgh]
            ])
benchMap ncgh n = env (do
            tMap <- testMap n n :: IO (Map A B) 
            tChanges <- testChanges n ncgh :: IO [MultiChange (IncMap.AtomicChange A B)]
            let tChangesScan = scanl (\x y -> x ++ [y]) [] tChanges
            let testValAndChanges = zip (replicate (ncgh+1) tMap) tChangesScan
            tFun  <- testFun n :: IO (B -> C)
            return (testValAndChanges,tFun))
         (\ ~(valAndChanges,fun) ->
            bgroup ("map " ++ show n) $ concat 
            [map (\x -> bench ("original change=" ++ show x) $ nf (checkOriginal (Map.map fun)) (valAndChanges !! x)) [0..ncgh],
             map (\x -> bench ("incremental change=" ++ show x) $ nf (checkTrans (IncMap.map fun)) (valAndChanges !! x)) [0..ncgh]
            ])

benchLookup ncgh n = env (do
            tMap <- testMap n n :: IO (Map A B) 
            tChanges <- testChanges n ncgh :: IO [MultiChange (IncMap.AtomicChange A B)]
            let tChangesScan = scanl (\x y -> x ++ [y]) [] tChanges
            let testValAndChanges = zip (replicate (ncgh+1) tMap) tChangesScan
            tKey  <- testKey n :: IO (A)
            return (testValAndChanges,tKey))
         (\ ~(valAndChanges,key) ->
            bgroup ("lookup " ++ show n) $ concat 
            [map (\x -> bench ("original change=" ++ show x) $ nf (checkOriginal (Map.lookup key)) (valAndChanges !! x)) [0..ncgh],
             map (\x -> bench ("incremental change=" ++ show x) $ nf (checkTrans (IncMap.lookup key)) (valAndChanges !! x)) [0..ncgh]
            ])
benchMember ncgh n = env (do
            tMap <- testMap n n :: IO (Map A B) 
            tChanges <- testChanges n ncgh :: IO [MultiChange (IncMap.AtomicChange A B)]
            let tChangesScan = scanl (\x y -> x ++ [y]) [] tChanges
            let testValAndChanges = zip (replicate (ncgh+1) tMap) tChangesScan
            tKey  <- testKey n :: IO (A)
            return (testValAndChanges,tKey))
         (\ ~(valAndChanges,key) ->
            bgroup ("member " ++ show n) $ concat 
            [map (\x -> bench ("original change=" ++ show x) $ nf (checkOriginal (Map.member key)) (valAndChanges !! x)) [0..ncgh],
             map (\x -> bench ("incremental change=" ++ show x) $ nf (checkTrans (IncMap.member key)) (valAndChanges !! x)) [0..ncgh]
            ])
benchPartition ncgh n = env (do
            tMap <- testMap n n :: IO (Map A B) 
            tChanges <- testChanges n ncgh :: IO [MultiChange (IncMap.AtomicChange A B)]
            let tChangesScan = scanl (\x y -> x ++ [y]) [] tChanges
            let testValAndChanges = zip (replicate (ncgh+1) tMap) tChangesScan
            tFun  <- testFun n :: IO (B -> Bool)
            return (testValAndChanges,tFun))
         (\ ~(valAndChanges,fun) ->
            bgroup ("partition " ++ show n) $ concat 
            [map (\x -> bench ("original change=" ++ show x) $ nf (checkOriginal (Map.partition fun)) (valAndChanges !! x)) [0..ncgh],
             map (\x -> bench ("incremental change=" ++ show x) $ nf (checkTrans (IncMap.partition fun)) (valAndChanges !! x)) [0..ncgh]
            ])

benchUnion ncgh n = env (do
            tMap1 <- testMap n n :: IO (Map A B) 
            tMap2 <- testMap n n :: IO (Map A B) 
            tChanges <- testChanges n ncgh :: IO [MultiChange (Tuple.AtomicChange (Map A B) (Map A B))]
            let tChangesScan = scanl (\x y -> x ++ [y]) [] tChanges
            let testValAndChanges = zip (replicate (ncgh+1) (tMap1,tMap2)) tChangesScan
            return testValAndChanges)
         (\ ~(valAndChanges) ->
            bgroup ("union " ++ show n) $ concat 
            [map (\x -> bench ("original change=" ++ show x) $ nf (checkOriginal (uncurry Map.union)) (valAndChanges !! x)) [0..ncgh],
             map (\x -> bench ("incremental change=" ++ show x) $ nf (checkTrans (IncMap.union)) (valAndChanges !! x)) [0..ncgh]
            ])

benchDifference ncgh n = env (do
            tMap1 <- testMap n n :: IO (Map A B) 
            tMap2 <- testMap n n :: IO (Map A B) 
            tChanges <- testChanges n ncgh :: IO [MultiChange (Tuple.AtomicChange (Map A B) (Map A B))]
            let tChangesScan = scanl (\x y -> x ++ [y]) [] tChanges
            let testValAndChanges = zip (replicate (ncgh+1) (tMap1,tMap2)) tChangesScan
            return testValAndChanges)
         (\ ~(valAndChanges) ->
            bgroup ("difference " ++ show n) $ concat 
            [map (\x -> bench ("original change=" ++ show x) $ nf (checkOriginal (uncurry Map.difference)) (valAndChanges !! x)) [0..ncgh],
             map (\x -> bench ("incremental change=" ++ show x) $ nf (checkTrans (IncMap.difference)) (valAndChanges !! x)) [0..ncgh]
            ])

benchIntersection ncgh n = env (do
            tMap1 <- testMap n n :: IO (Map A B) 
            tMap2 <- testMap n n :: IO (Map A B) 
            tChanges <- testChanges n ncgh :: IO [MultiChange (Tuple.AtomicChange (Map A B) (Map A B))]
            let tChangesScan = scanl (\x y -> x ++ [y]) [] tChanges
            let testValAndChanges = zip (replicate (ncgh+1) (tMap1,tMap2)) tChangesScan
            return testValAndChanges)
         (\ ~(valAndChanges) ->
            bgroup ("intersection " ++ show n) $ concat 
            [map (\x -> bench ("original change=" ++ show x) $ nf (checkOriginal (uncurry Map.intersection)) (valAndChanges !! x)) [0..ncgh],
             map (\x -> bench ("incremental change=" ++ show x) $ nf (checkTrans (IncMap.intersection)) (valAndChanges !! x)) [0..ncgh]
            ])

benchKeysSet ncgh n = env (do
            tMap <- testMap n n :: IO (Map A B) 
            tChanges <- testChanges n ncgh :: IO [MultiChange (IncMap.AtomicChange A B)]
            let tChangesScan = scanl (\x y -> x ++ [y]) [] tChanges
            let testValAndChanges = zip (replicate (ncgh+1) tMap) tChangesScan
            return testValAndChanges)
         (\ ~(valAndChanges) ->
            bgroup ("keysSet " ++ show n) $ concat 
            [map (\x -> bench ("original change=" ++ show x) $ nf (checkOriginal (Map.keysSet)) (valAndChanges !! x)) [0..ncgh],
             map (\x -> bench ("incremental change=" ++ show x) $ nf (checkTrans (IncMap.keysSet)) (valAndChanges !! x)) [0..ncgh]
            ])

benchSubmapOf ncgh n = env (do
            tMap1 <- testMap n n :: IO (Map A B) 
            tMap2 <- testMap n n :: IO (Map A B) 
            tChanges <- testChanges n ncgh :: IO [MultiChange (Tuple.AtomicChange (Map A B) (Map A B))]
            let tChangesScan = scanl (\x y -> x ++ [y]) [] tChanges
            let testValAndChanges = zip (replicate (ncgh+1) (tMap1,tMap1)) tChangesScan
            return testValAndChanges)
         (\ ~(valAndChanges) ->
            bgroup ("submapOf " ++ show n) $ concat 
            [map (\x -> bench ("original change=" ++ show x) $ nf (checkOriginal (uncurry Map.isSubmapOf)) (valAndChanges !! x)) [0..ncgh],
             map (\x -> bench ("incremental change=" ++ show x) $ nf (checkTrans (IncMap.isSubmapOf)) (valAndChanges !! x)) [0..ncgh]
            ])

benchSplit ncgh n = env (do
            tMap <- testMap n n :: IO (Map A B) 
            tKey <- testKey n :: IO (A)
            tChanges <- testChanges n ncgh :: IO [MultiChange (IncMap.AtomicChange A B)]
            let tChangesScan = scanl (\x y -> x ++ [y]) [] tChanges
            let testValAndChanges = zip (replicate (ncgh+1) tMap) tChangesScan
            return (testValAndChanges,tKey))
         (\ ~(valAndChanges,key) ->
            bgroup ("split " ++ show n) $ concat 
            [map (\x -> bench ("original change=" ++ show x) $ nf (checkOriginal (Map.split key)) (valAndChanges !! x)) [0..ncgh],
             map (\x -> bench ("incremental change=" ++ show x) $ nf (checkTrans (IncMap.split key)) (valAndChanges !! x)) [0..ncgh]
            ])

benchMapKeys ncgh n = env (do
            tMap <- testMap n n :: IO (Map A B) 
            tFun <- testFun n :: IO (A -> C)
            tChanges <- testChanges n ncgh :: IO [MultiChange (IncMap.AtomicChange A B)]
            let tChangesScan = scanl (\x y -> x ++ [y]) [] tChanges
            let testValAndChanges = zip (replicate (ncgh+1) tMap) tChangesScan
            return (testValAndChanges,tFun))
         (\ ~(valAndChanges,fun) ->
            bgroup ("mapKeys " ++ show n) $ concat 
            [map (\x -> bench ("original change=" ++ show x) $ nf (checkOriginal (Map.mapKeys fun)) (valAndChanges !! x)) [0..ncgh],
             map (\x -> bench ("incremental change=" ++ show x) $ nf (checkTrans (IncMap.mapKeys fun)) (valAndChanges !! x)) [0..ncgh]
            ])