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.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


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 "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