module RandList
  ( RandList
  , empty
  , isEmpty
  , cons
  , head
  , tail
  , lookup
  , update
  ) where

import           Prelude hiding (head, lookup, tail)

data BLT a = Leaf a
           | Fork !Int (BLT a) (BLT a)
   deriving (Eq, Show)

data Digit a = Zero
             | One (BLT a)
   deriving (Eq, Show)

type RandList a = [Digit a]

empty :: RandList a
empty = []

isEmpty :: RandList a -> Bool
isEmpty = null

cons :: a -> RandList a -> RandList a
cons x = insTree (Leaf x)

head :: RandList a -> a
head ts = x where (Leaf x, _) = borrowTree ts

tail :: RandList a -> RandList a
tail ts = ts' where (_, ts') = borrowTree ts

lookup :: RandList a -> Int -> a
lookup [] _ = error "Bad Index"
lookup (Zero : ts) i = lookup ts i
lookup (One t : ts) i
   | i < size t = lookupTree t i
   | otherwise  = lookup ts (i - size t)

update :: RandList a -> Int -> a -> RandList a
update [] _ _ = error "Bad Index"
update (Zero : ts) i y = Zero : update ts i y
update (One t : ts) i y
   | i < size t = One (updateTree t i y) : ts
   | otherwise  = One t : update ts (i - size t) y

-- Utility functions

size :: BLT a -> Int
size (Leaf _)      = 1
size (Fork n _ _ ) = n

link :: BLT a -> BLT a -> BLT a
link t1 t2 = Fork (size t1 + size t2) t1 t2

insTree :: BLT a -> RandList a -> RandList a
insTree t [] = [One t]
insTree t (Zero : ts) = One t : ts
insTree t1 (One t2 : ts) = Zero : insTree (link t1 t2) ts

borrowTree :: RandList a -> (BLT a, RandList a)
borrowTree [One t]      = (t, [])
borrowTree (One t : ts) = (t, Zero : ts)
borrowTree (Zero  : ts) = (t1, One t2 : ts')
  where (Fork _ t1 t2, ts') = borrowTree ts

lookupTree :: BLT a -> Int -> a
lookupTree (Leaf x) 0 = x
lookupTree (Leaf _) _ = error "Bad Index"
lookupTree (Fork n t1 t2) i
   | i < (n `div` 2) = lookupTree t1 i
   | otherwise       = lookupTree t2 (i - n `div` 2)

updateTree :: BLT a -> Int -> a -> BLT a
updateTree (Leaf _) 0 y = Leaf y
updateTree (Leaf _) _ _ = error "Bad Index"
updateTree (Fork n t1 t2) i y
   | i < (n `div` 2) = Fork n (updateTree t1 i y) t2
   | otherwise       = Fork n t1 (updateTree t2 (i - n `div` 2) y)