content: module AVLTree where 
  
  -- from: https://gist.github.com/gerard/109729
  
  data BT a = L 
            | N { dep   :: !Int 
                , value :: a   
                , left  :: !(BT a)
                , right :: !(BT a) 
                }
            deriving (Show, Eq)
  
  {-@ data BT a [dep] = L 
                      | N (dep   :: !Int) 
                          (value :: a   )
                          (left  :: Bt a) 
                          (right :: BT a) 
    @-}
  
  -- 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 d _ _ _) = d 
  
  {-@ measure depth   :: BT a -> Int 
      depth (L)         = 0
      depth (N d x l r) = 1 + (if (depth l) <= (depth r) then (depth r) else (depth l)) 
    @-}
  
  inorder             :: BT a -> [a]
  inorder L           = []
  inorder (N _ v l r) = inorder l ++ [v] ++ inorder r
   
  -- 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
  