module Main where
import qualified Data.Map as Map
import qualified Data.IntMap as IntMap
import Data.Array.IArray
import Data.Array.Diff
import Data.Array.Unboxed
import Data.Array.ST
import Data.LazyArray
import Data.LazyArray.Lowlevel
import Data.Maybe
import CPUTime

import Debug.Trace

graph_complete::Int->Array Int [Int]
graph_complete n = array (1,n) [(i,[j|j<-[1..n],j/=i]) | i<-[1..n]]
graph_narrow::Int->Array Int [Int]
graph_narrow n = array (1,n) [(i,[j|j<-[(i-10)..(i+10)],j>=1,j<=n,j/=i]) | i<-[1..n]]
graph_path::Int->Array Int [Int]
graph_path n = array (1,n) [(i,if i==n then [i-1] else if i==1 then [i+1] else [i-1,i+1]) | i<-[1..n]]

-- all these methods are not generalised to be as fast as possible
dnum_la::((Int,Int)->[(Int,Int)]->Array Int Int)->Array Int [Int]->Int->Array Int Int
dnum_la create g s = marks where
  list = dfs' [s] 0
  marks = create (bounds g) list
  dfs' []     _ = []
  dfs' (s:ss) n = (s,n) : if n == marks!s then dfs' ((g!s)++ss) (n+1) else dfs' ss n

dnum_la_maybe::((Int,Int)->[(Int,Int)]->Array Int (Maybe Int))->Array Int [Int]->Int->Array Int (Maybe Int)
dnum_la_maybe create g s = marks where
  list = dfs' [s] 0
  marks = create (bounds g) list
  dfs' []     _ = []
  dfs' (s:ss) n = (s,n) : if n == (case marks!s of Just e->e) then dfs' ((g!s)++ss) (n+1) else dfs' ss n

dnum_la_lowlevel::Array Int [Int]->Int->Array Int Int
dnum_la_lowlevel g s = laFreeze marks where
  list = dfs' [s] 0
  marks = laCreate (-1) (bounds g) list
  dfs' []     _ = []
  dfs' (s:ss) n = (s,n) : if n == (marks `laAt` s) then dfs' ((g!s)++ss) (n+1) else dfs' ss n

dnum_la_lowlevel_maybe::Array Int [Int]->Int->Array Int (Maybe Int)
dnum_la_lowlevel_maybe g s = laFreeze marks where
  list = dfs' [s] 0
  marks = mlaCreate (bounds g) list
  dfs' []     _ = []
  dfs' (s:ss) n = (s,n) : if n == case (marks `laAt` s) of Just e->e then dfs' ((g!s)++ss) (n+1) else dfs' ss n

dnum_diff g s = dfs' [s] (array (bounds g) [(i,-1)|i<-(range (bounds g))]) 0 where
  dfs'::[Int]->DiffArray Int Int->Int->DiffArray Int Int
  dfs' []     m _ = m
  dfs' (s:ss) m n = if m!s==(-1) then dfs' ((g!s)++ss) (m//[(s,n)]) (n+1) else dfs' ss m n

dnum_map::Array Int [Int]->Int->[(Int,Int)]
dnum_map g s = dfs' [s] Map.empty 0 where
  dfs' []     m _ = Map.assocs m
  dfs' (s:ss) m n = if Map.notMember s m then dfs' ((g!s)++ss) (Map.insert s n m) (n+1) else dfs' ss m n

dnum_imap::Array Int [Int]->Int->[(Int,Int)]
dnum_imap g s = dfs' [s] IntMap.empty 0 where
  dfs' []     m _ = IntMap.assocs m
  dfs' (s:ss) m n = if IntMap.notMember s m then dfs' ((g!s)++ss) (IntMap.insert s n m) (n+1) else dfs' ss m n

dnum_st::Array Int [Int]->Int->Array Int Int
dnum_st g s = runSTArray $ newArray (bounds g) (-1) >>= dfs' [s] 0 where
  dfs' []     _ m = return m
  dfs' (s:ss) n m = do sn<-readArray m s
                       if sn==(-1) then writeArray m s n >> dfs' ((g!s)++ss) (n+1) m else dfs' ss n m

dnum_stu::Array Int [Int]->Int->UArray Int Int
dnum_stu g s = runSTUArray $ newArray (bounds g) (-1) >>= dfs' [s] 0 where
  dfs' []     _ m = return m
  dfs' (s:ss) n m = do sn<-readArray m s
                       if sn==(-1) then writeArray m s n >> dfs' ((g!s)++ss) (n+1) m else dfs' ss n m

touch_array arr = foldl1 (+) (tail $ elems arr)
touch_list list = foldl (\x (i,e)->x+e) 0 list
touch_array_maybe arr = foldl (\x (Just y)->x+y) 0 (tail $ elems arr)
touch_array_f f arr = foldl (\x y->x+f y) 0 (tail $ elems arr)

graphs = [gp, isol gp, gn, isol gn, gc, isol gc]
 where
  gp=graph_path   10000
  gn=graph_narrow  2500
  gc=graph_complete 250
  isol g = array (0,snd (bounds g)) $ (0,[]):assocs g

measure str f touch = do putStr str; putStr ": "; 
                         foldr1 (>>) $ map measure' graphs
                         putStrLn ""
 where
  measure' g = do start<-getCPUTime
                  end<-foldl1 seq [touch (f g i) | i<-[1..100]] `seq` getCPUTime
                  putStr $ show $ (end-start) `div` 1000000000; putStr "; "

main = do
  foldl1 seq (map (touch_array_f length) graphs) `seq` putStrLn "Testing..."
  measure "lArrayFirst"                 (dnum_la (lArrayFirst (-1)))  touch_array
  measure "lArrayFirst using lArrayMap" (dnum_la (lArrayMap (\x->case x of []->(-1);e:_->e))) touch_array
  measure "lArrayMaybe"                 (dnum_la_maybe lArrayMaybe)   touch_array_maybe
  measure "lArrayMaybe using lArrayMap" (dnum_la_maybe (lArrayMap listToMaybe)) touch_array_maybe
  measure "laCreate"                     dnum_la_lowlevel             touch_array
  measure "mlaCreate"                    dnum_la_lowlevel_maybe       touch_array_maybe
  measure "DiffArray"                    dnum_diff                    touch_array
  measure "Map"                          dnum_map                     touch_list
  measure "IntMap"                       dnum_imap                    touch_list
  measure "runSTArray"                   dnum_st                      touch_array
  measure "runSTUArray"                  dnum_stu                     touch_array
  measure "lArrayFirst"                 (dnum_la (lArrayFirst (-1)))  touch_array
  measure "lArrayFirst using lArrayMap" (dnum_la (lArrayMap (\x->case x of []->(-1);e:_->e))) touch_array
  measure "lArrayMaybe"                 (dnum_la_maybe lArrayMaybe)   touch_array_maybe
  measure "lArrayMaybe using lArrayMap" (dnum_la_maybe (lArrayMap listToMaybe)) touch_array_maybe
  measure "laCreate"                     dnum_la_lowlevel             touch_array
  measure "mlaCreate"                    dnum_la_lowlevel_maybe       touch_array_maybe
  measure "DiffArray"                    dnum_diff                    touch_array
  measure "Map"                          dnum_map                     touch_list
  measure "IntMap"                       dnum_imap                    touch_list
  measure "runSTArray"                   dnum_st                      touch_array
  measure "runSTUArray"                  dnum_stu                     touch_array