import qualified Data.IntMap as Map
import Data.List as List
import Data.Maybe (maybe)
import MersenneTwister

type Map = Map.IntMap Int

n    = 1000000  --size of map
kmax = 10000000 --largest key

makeRands n s = take n $ map (fromIntegral.(`mod`kmax)) (mersenneTwister s)

sequ1 = [0..n]
sequ2 = [n+1..2*n]
sequ3 = [2*n+1..3*n]

--maps with random keys
rmap1 = Map.fromList . take n $ makeRands n 42 `zip` [0..]
rmap2 = Map.fromList . take n $ makeRands n 44 `zip` [0..]
rmap3 = Map.fromList . take n $ makeRands n 46 `zip` [0..]

--maps with random, duplicate keys
dmap1 = Map.fromList . take n $ (map (`div`1000) $ makeRands n 73) `zip` [0..]
dmap2 = Map.fromList . take n $ (map (`div`1000) $ makeRands n 75) `zip` [0..]
dmap3 = Map.fromList . take n $ (map (`div`1000) $ makeRands n 77) `zip` [0..]

--maps with sequential, non-overlapping keys
smap1 = Map.fromAscList $ sequ1 `zip` [0..]
smap2 = Map.fromAscList $ sequ2 `zip` [0..]
smap3 = Map.fromAscList $ sequ3 `zip` [0..]

manyLittle :: [Map]
manyLittle = 
  [ Map.fromList . take n $ makeRands 100 i `zip` [0..] | i<-[0..] ]

-- intersection of many maps
bigCap :: [Map] -> Map
bigCap = List.foldl1' (\a b -> Map.difference a b) 



--query1 : sum all keys
query1 :: Map -> Int
query1 = List.foldl' (+) 0 . Map.keys

--query2 : sum a random subset of values
query2 :: Int -> Map -> Int
query2 n m =
  let f s k = seq s $ maybe s (+s) (Map.lookup k m)
  in  List.foldl' f 0 (makeRands n 2001)

--query3 : union maps and compute key sum
query3 :: [Map] -> Int
query3 = query1 . Map.unions

--query4 : intersect maps and compute key sum
query4 :: [Map] -> Int
query4 = query1 . bigCap

--query5 : union maps and compute sum of subset of values
query5 :: Int -> [Map] -> Int
query5 n = query2 n . Map.unions

--query6 : intersection of maps and compute sum of subset of values
query6 :: Int -> [Map] -> Int
query6 n = query2 n . bigCap

query7 :: [Map] -> Int
query7 (x:xs) = query1 $ List.foldl' Map.difference x xs



main = 
  do
  print $ query5 1 (take 10000 manyLittle)