module AVLTree where 
data BT = L | N Int BT BT deriving (Show, Eq)

-- Nice, small and useful functions
empty = L

-- You could say, depth L == Nothing depth (N v L L) == Just 0, but works for
-- me better this way:
depth L = 0
depth (N _ t u) = (max (depth t) (depth u)) + 1

inorder :: BT -> [Int]
inorder L = []
inorder (N v t u) = inorder t ++ [v] ++ inorder u

left (N _ t _) = t
right (N _ _ u) = u
value (N v _ _) = v

btmin = head . inorder

-- FIXME: Could be cleaner BT -> Int using left and right of BT
balFactor :: BT -> BT -> Int
balFactor t u = (depth t) - (depth u)

-- Tricky but easy: we return a binary list with the route to the node
search :: BT -> Int -> Maybe [Int]
search L s = Nothing
search (N v t u) s 
    | v == s                  = Just []
    | (search t s) /= Nothing = fmap ((:) 0) (search t s)
    | (search u s) /= Nothing = fmap ((:) 1) (search u s)
    | otherwise               = Nothing

-- Complementary to search: get the node with the path
getelem :: BT -> [Int] -> Maybe Int
getelem L _ = Nothing
getelem (N v _ _) [] = Just v
getelem (N v t u) (x:xs)
    | x == 0    = getelem t xs
    | otherwise = getelem u xs

-- If you get confused (I do), check this nice picture:
-- http://en.wikipedia.org/wiki/Image:Tree_Rebalancing.gif
balanceLL (N v (N vl tl ul) u)              = (N vl tl (N v ul u))
balanceLR (N v (N vl tl (N vlr tlr ulr)) u) = (N vlr (N vl tl tlr) (N v ulr u))
balanceRL (N v t (N vr (N vrl trl url) ur)) = (N vrl (N v t trl) (N vr url ur)) 
balanceRR (N v t (N vr tr ur))              = (N vr (N v t tr) ur)

-- Balanced insert
insert :: BT -> Int -> BT
insert L i = (N i L L)
insert (N v t u) i
    | i == v = (N v t u)
    | i < v && (balFactor ti u) ==  2 && i < value t = balanceLL (N v ti u)
    | i < v && (balFactor ti u) ==  2 && i > value t = balanceLR (N v ti u)
    | i > v && (balFactor t ui) == -2 && i < value u = balanceRL (N v t ui)
    | i > v && (balFactor t ui) == -2 && i > value u = balanceRR (N v t ui)
    | i < v  = (N v ti u)
    | i > v  = (N v t ui)
        where ti = insert t i
              ui = insert u i

-- Balanced delete
delete :: BT -> Int -> BT
delete L d = L
delete (N v L L) d = if v == d then L else (N v L L)
delete (N v t L) d = if v == d then t else (N v t L)
delete (N v L u) d = if v == d then u else (N v L u)
delete (N v t u) d
    | v == d                            = (N mu t dmin)
    | v > d && abs (balFactor dt u) < 2 = (N v dt u)
    | v < d && abs (balFactor t du) < 2 = (N v t du)
    | v > d && (balFactor (left u) (right u)) < 0 = balanceRR (N v dt u) 
    | v < d && (balFactor (left t) (right t)) > 0 = balanceLL (N v t du)
    | v > d                                       = balanceRL (N v dt u)
    | v < d                                       = balanceLR (N v t du)
        where dmin = delete u mu
              dt   = delete t d
              du   = delete u d
              mu   = btmin u

-- Test Functions
load :: BT -> [Int] -> BT
load t []     = t
load t (x:xs) = insert (load t xs) x

unload :: BT -> [Int] -> BT
unload t []     = t
unload t (x:xs) = delete (unload t xs) x

sort :: [Int] -> [Int]
sort = inorder . (load empty)

isBalanced L = True
isBalanced (N _ t u) = isBalanced t && isBalanced u && abs (balFactor t u) < 2