module Main(main) where

import Prelude hiding(lookup)
import System.Random
import System.CPUTime
import Data.FiniteMap
import GHC.Prim

data Trie k v = Trie !(Maybe v) !(TrieItem k v)

data TrieItem k v
  = Empty
  | Branch 
      !k
      ![k]
      !(Maybe v)
      Int#
      !(TrieItem k v)
      !(TrieItem k v)
      !(TrieItem k v)

{--------------------------------------------------------------------
  Query
--------------------------------------------------------------------}

null :: Trie k v -> Bool
null (Trie Nothing Empty) = True
null _                    = False

size :: Trie k v -> Int
size (Trie mb_v item) = mb_inc mb_v (itemSize item)
  where
    mb_inc Nothing s = s
    mb_inc _       s = s+1
    
    itemSize Empty = 0
    itemSize (Branch _ _ mb_v _ fm_l fm_c fm_r) =
      mb_inc mb_v (itemSize fm_l + itemSize fm_c + itemSize fm_r)

lookup :: Ord k => [k] -> Trie k v -> Maybe v
lookup []      (Trie mb_v item) = mb_v
lookup (k:key) (Trie mb_v item) = lookupItem item k key
  where
    lookupItem Empty                                   k key = Nothing
    lookupItem (Branch p prefix mb_v _ fm_l fm_c fm_r) k key =
      case compare k p of
        LT -> lookupItem fm_l k key
        EQ -> checkPrefix key prefix
        GT -> lookupItem fm_r k key
      where
        checkPrefix []      []         = mb_v
        checkPrefix []      (p:prefix) = Nothing
        checkPrefix (k:key) []         = lookupItem fm_c k key
        checkPrefix (k:key) (p:prefix)
          | k == p = checkPrefix key prefix

{--------------------------------------------------------------------
  Construction
--------------------------------------------------------------------}

empty :: Trie k v
empty = Trie Nothing Empty

singleton :: [k] -> v -> Trie k v
singleton []      v = Trie (Just v) Empty
singleton (k:key) v = Trie Nothing  (Branch k key (Just v) 1# Empty Empty Empty)

{--------------------------------------------------------------------
  Insertion
  [insert] is the inlined version of [insertWith (\x y -> x)]
--------------------------------------------------------------------}

insert :: Ord k => Trie k v -> [k] -> v -> Trie k v
insert = insertWith (\new old -> new)

insertWith :: Ord k => (a -> a -> a)
         -> Trie k a -> [k] -> a -> Trie k a
insertWith combiner (Trie mb_v item) [] v =
  Trie (combineValues combiner mb_v v) item
insertWith combiner (Trie mb_v item) (k:key) v =
  Trie mb_v (insertItem item k key)
  where
    insertItem Empty                                       k key =
      Branch k key (Just v) 1# Empty Empty Empty
    insertItem (Branch p prefix mb_v size# fm_l fm_c fm_r) k key =
      case compare k p of
        LT -> mkBalBranch p prefix mb_v (insertItem fm_l k key) fm_c fm_r
        EQ -> checkPrefix [] key prefix
        GT -> mkBalBranch p prefix mb_v fm_l fm_c (insertItem fm_r k key)
      where
        checkPrefix cs []        []           =
          Branch p prefix (combineValues combiner mb_v v) size# fm_l fm_c fm_r
        checkPrefix cs []        (p1:prefix1) =
          Branch p cs' (Just v) size# fm_l fm_c' fm_r
          where
            cs'   = reverse cs
            fm_c' = Branch p1 prefix1 mb_v 1# Empty fm_c Empty
        checkPrefix cs (k1:key1) []           =
          Branch p prefix mb_v size# fm_l (insertItem fm_c k1 key1) fm_r
        checkPrefix cs (k1:key1) (p1:prefix1)
          | k1 == p1  = checkPrefix (k1:cs) key1 prefix1
          | otherwise =
              Branch p cs' Nothing size# fm_l fm_c' fm_r
              where
                cs'   = reverse cs
                fm_c' = insertItem (Branch p1 prefix1 mb_v 1# Empty fm_c Empty) k1 key1

-- ---------------------------------------------------------------------------
-- The implementation of balancing

-- Basic construction of a @TrieItem@:

-- @mkBranch@ simply gets the size component right.  This is the ONLY
-- (non-trivial) place the Branch object is built, so the ASSERTion
-- recursively checks consistency.

-- sIZE_RATIO :: Int#
#define sIZE_RATIO 5#

sizeTI Empty		             = 0#
sizeTI (Branch _ _ _ size _ _ _) = size

mkBranch :: k
         -> [k]
         -> Maybe v
         -> TrieItem k v 
         -> TrieItem k v
         -> TrieItem k v
         -> TrieItem k v
mkBranch k key mb_v fm_l fm_c fm_r =
  Branch k key mb_v (1# +# sizeTI fm_l +# sizeTI fm_r) fm_l fm_c fm_r

-- ---------------------------------------------------------------------------
-- Balanced construction of a @TrieItem@

-- @mkBalBranch@ rebalances, assuming that the subtrees aren't too far
-- out of whack.

mkBalBranch :: Ord k
            => k 
            -> [k]
            -> Maybe v
            -> TrieItem k v 
            -> TrieItem k v
            -> TrieItem k v
            -> TrieItem k v
mkBalBranch k key mb_v fm_L fm_C fm_R
  | size_l +# size_r <# 2# =
      mkBranch k key mb_v fm_L fm_C fm_R
  | size_r ># (sIZE_RATIO *# size_l) =	-- Right tree too big
      case fm_R of
	    Branch _ _ _ _ fm_rl fm_rc fm_rr
		  | sizeTI fm_rl <# (2# *# sizeTI fm_rr) -> single_L fm_L fm_C fm_R
		  | otherwise	   	                     -> double_L fm_L fm_C fm_R
	      -- Other case impossible
  | size_l ># (sIZE_RATIO *# size_r) =	-- Left tree too big
      case fm_L of
	    Branch _ _ _ _ fm_ll fm_lc fm_lr
		  | sizeTI fm_lr <# (2# *# sizeTI fm_ll) -> single_R fm_L fm_C fm_R
		  | otherwise		                     -> double_R fm_L fm_C fm_R
	      -- Other case impossible
  | otherwise =				-- No imbalance
      mkBranch k key mb_v fm_L fm_C fm_R
  where
    size_l   = sizeTI fm_L
    size_r   = sizeTI fm_R

    single_L fm_l fm_c (Branch k_r key_r mb_v_r _ fm_rl fm_rc fm_rr) =
	  mkBranch k_r key_r mb_v_r (mkBranch k key mb_v fm_l fm_c fm_rl) fm_rc fm_rr

    double_L fm_l fm_c (Branch k_r key_r mb_v_r _ (Branch k_rl key_rl mb_v_rl _ fm_rll fm_rlc fm_rlr) fm_rc fm_rr) =
	  mkBranch k_rl key_rl mb_v_rl (mkBranch k key mb_v fm_l fm_c fm_rll) fm_rlc
				 (mkBranch k_r key_r mb_v_r fm_rlr fm_rc fm_rr)

    single_R (Branch k_l key_l mb_v_l _ fm_ll fm_lc fm_lr) fm_c fm_r =
	  mkBranch k_l key_l mb_v_l fm_ll fm_lc (mkBranch k key mb_v fm_lr fm_c fm_r)

    double_R (Branch k_l key_l mb_v_l _ fm_ll fm_lc (Branch k_lr key_lr mb_v_lr _ fm_lrl fm_lrc fm_lrr)) fm_c fm_r =
	  mkBranch k_lr key_lr mb_v_lr (mkBranch k_l key_l mb_v_l fm_ll fm_lc fm_lrl) fm_lrc
				 (mkBranch k key mb_v fm_lrr fm_c fm_r)

combineValues :: (a -> a -> a) -> Maybe a -> a -> Maybe a
combineValues combiner Nothing  y = Just y
combineValues combiner (Just x) y = Just (combiner y x)

main = do
  s <- genStrings
  let strings = take 100000 s
      trie = foldr (\s t -> insertWith (+) t s 1) empty strings
      fm = foldr (\s t -> addToFM_C (+) t s 1) emptyFM strings
  print strings
  t1 <- getCPUTime
  trie `seq` return ()
  t2 <- getCPUTime
  print (gg (\s -> lookup s trie) strings 0)
  t3 <- getCPUTime
  fm `seq` return ()
  t4 <- getCPUTime
  print (gg (lookupFM fm) strings 0)
  t5 <- getCPUTime
  print (t2-t1,t3-t2)
  print (t4-t3,t5-t4)
  let percDiff :: Integer -> Integer -> Double
      percDiff x y = (1-(fromIntegral x)/(fromIntegral y))*100
  print (percDiff (t2-t1) (t4-t3),percDiff (t3-t2) (t5-t4))

gg :: (String -> Maybe Int) -> [String] -> Int -> Int
gg f []         s = s
gg f (key:keys) s =
  case f key of
    Just s' -> gg f keys (s+s')
    Nothing -> gg f keys s

genStrings :: IO [String]
genStrings = do
  g <- getStdGen
  return (loop g)
  where
    loop g =
      case gen g of
        (s,g) -> s:loop g

    gen g =
      case randomR (0, 50::Int) g of
        (n,  g) | n > 40    -> ([],g)
                | otherwise ->
           case randomR ('A','Z') g of
             (c,g) -> 
                case gen g of
                   (cs,g) -> (c:cs,g)