{-# LANGUAGE RankNTypes #-}

import Criterion.Main
import Data.List
import Data.Functor
import Data.Hashable
import qualified Data.ByteString.Char8 as BS
import qualified Data.ByteString as BSC

import qualified Data.HashMap.Strict as HashMap (HashMap, insertWith, size, empty)
import qualified Data.HashMap.Lazy as LazyHashMap (HashMap, insertWith, size, empty)

import qualified Data.Map.Strict as TreeMap (Map, insertWith, size, empty)
import qualified Data.Map.Lazy as LazyTreeMap (Map, insertWith, size, empty)

import qualified Data.Trie as TrieMap (Trie, size, empty)
import qualified Data.Trie.Convenience as TrieMap (insertWith)

import qualified Data.Discrimination.Grouping as EKmett (nub)

data MapImpl m k a = MapImpl {
	insertWith :: (Ord k, Hashable k) => (a -> a -> a) -> k -> a -> m -> m,
	size :: m -> Int,
	empty :: m
}

update impl str m = (insertWith impl) (flip const) str ((size impl) m) m
computeMap impl = foldr (update impl) (empty impl)
test impl list = (size impl) (computeMap impl list)

main = do
	list <- map BSC.unpack . filter (not . BS.null) . BS.split '\n' <$> BS.getContents
	defaultMain [ bench "hashMap"      $ whnf (test hashMap) list
	            , bench "lazyHashMap"  $ whnf (test lazyHashMap) list
	            , bench "treeMap"      $ whnf (test treeMap) list
	            , bench "lazyTreeMap"  $ whnf (test lazyTreeMap) list
	            --, bench "trieMap"      $ whnf (test trieMap) list
	            , bench "ekmett nub"   $ whnf (length . EKmett.nub) list
	            , bench "group . sort" $ whnf (length . group . sort) list
	            , bench "nub"          $ whnf (length . nub) list
	            ]

--main = do
--	list <- filter (not . BS.null) . BS.split '\n' <$> BS.getContents
--	print $ test hashMap list
--	print $ test lazyHashMap list
--	print $ test treeMap list
--	print $ test lazyTreeMap list
--	print $ test trieMap list

hashMap = MapImpl {
	insertWith = HashMap.insertWith,
	size = HashMap.size,
	empty = HashMap.empty
}

lazyHashMap = MapImpl {
	insertWith = LazyHashMap.insertWith,
	size = LazyHashMap.size,
	empty = LazyHashMap.empty
}

treeMap = MapImpl {
	insertWith = TreeMap.insertWith,
	size = TreeMap.size,
	empty = TreeMap.empty
}

lazyTreeMap = MapImpl {
	insertWith = LazyTreeMap.insertWith,
	size = LazyTreeMap.size,
	empty = LazyTreeMap.empty
}

trieMap = MapImpl {
	insertWith = TrieMap.insertWith,
	size = TrieMap.size,
	empty = TrieMap.empty
}