(* AVL Balanced TST *)

type 'a t = Node of ('a list * char * 'a t * 'a t * 'a t * int) | Empty

let empty = Empty

let mk_leaf value = Node ([value], '$', Empty, Empty, Empty, 0) (* $ is bogus *)

let rec find str tst =
  let last = (String.length str) - 1 in
  let rec find tst pos = 
    if pos > last then
      match tst with
          Node (x::_, _, _, _, _, _) -> x
        | _ -> raise Not_found
    else
      let cs = str.[pos]
      in match tst with
          Node (_, ck, left, _,  _, _) when cs < ck -> find left pos
        | Node (_, ck, _, same,  _, _) when cs = ck -> find same (pos+1)
        | Node (_, ck, _, _, right, _) when cs > ck -> find right pos
        | _ -> raise Not_found
  in find tst 0

let mem str trie =
  try ignore(find str trie); true
  with Not_found -> false

let rec mk_chain_rev str pos first acc =
  if pos < first then acc
  else mk_chain_rev str (pos-1) first
    (Node ([], str.[pos], Empty, acc, Empty, 1))

let height = function
    Empty -> 0
  | Node(_, _, _, _, _, h) -> h

let mk_node vals ck left same right =
  let hl = height left and hr = height right in
  let h = if hl >= hr then hl + 1 else hr + 1 in
    Node (vals, ck, left, same, right, h)

let bft = 1 (* 1 seems slightly better than 2 despite INRIA using 2 for Map *)
let debug = false

let bal vals ck left same right = 
  let hl = height left and hr = height right in
    if hl > hr + bft then
      match left with
          Empty -> invalid_arg "bal"
        | Node (l_vals, l_ck, l_left, l_same, l_right, _) ->
            if height l_left >= height l_right then (
(*               print_endline "Single Left"; *)
              mk_node l_vals l_ck l_left l_same 
                (mk_node vals ck l_right same right))
            else
              match l_right with
                  Empty -> invalid_arg "bal"
                | Node (lr_vals, lr_ck, lr_left, lr_same, lr_right, _) -> (
(*                     print_endline "Double Left"; *)
                    mk_node lr_vals lr_ck
                      (mk_node l_vals l_ck l_left l_same lr_left)
                      lr_same
                      (mk_node vals ck lr_right same right))
    else if hr > hl + bft then
      match right with
          Empty -> invalid_arg "bal"
        | Node (r_vals, r_ck, r_left, r_same, r_right, _) ->
            if height r_right >= height r_left then (
(*               print_endline "Single Right"; *)
              mk_node r_vals r_ck 
                (mk_node vals ck left same r_left) r_same r_right)
            else
              match r_left with
                  Empty -> invalid_arg "bal"
                | Node (rl_vals, rl_ck, rl_left, rl_same, rl_right, _) -> (
(*                     print_endline "Double Right"; *)
                    mk_node rl_vals rl_ck
                      (mk_node vals ck left same rl_left)
                      rl_same
                      (mk_node r_vals r_ck rl_right r_same r_right))
    else
      let h = if hl >= hr then hl + 1 else hr + 1
      in (
(*         print_endline "Balanced"; *)
        Node (vals, ck, left, same, right, h))


let add str value tst =
  let last = (String.length str) - 1 in
  let rec add tst pos =
    if pos > last then
      match tst with
          Empty -> mk_leaf value
        | Node (vals, ck, l, s, r, h) -> Node (value::vals, ck, l, s, r, h)
    else
      let cs = str.[pos]
      in match tst with
          Node (vals, ck, left, same, right, h) ->
            if cs < ck then bal vals ck (add left pos) same right
            else if cs > ck then bal vals ck left same (add right pos)
            else Node (vals, ck, left, (add same (pos+1)), right, h)
        | Empty -> mk_chain_rev str last pos (mk_leaf value)
  in
    add tst 0


(* let myt1 = Node ([], 'a', Empty, (mk_leaf 42) , Empty) (* a => 42 *) *)
(* let myt2 = Node ([], 'g', Empty, myt1, Empty) (* ga => 42 *) *)
(* let myt1 = Node ([1], 'a', Empty, (mk_leaf 99) , Empty) (* "" => 1 a => 99 *) *)
(* let myt2 = Node ([], 'v', myt2, myt1, Empty);; (* v => 1 ga => 42 va => 99 *) *)

(* find "" myt2;; *)
(* find "v" myt2;; *)
(* find "va" myt2;; *)
(* find "g" myt2;; *)
(* find "ga" myt2;; *)
(* find "ga$" myt2;; *)
(* find "gaz" myt2;; *)

(* mk_chain_rev "foo" 2 0 empty;; *)

(* let myt = add "foo" 42 empty;; *)
(* let myt = add "bar" 47 myt;; *)
(* let myt = add "zap" 99 myt;; *)
(* let myt = add "foo" 77 myt;; *)

(* find "foo" myt;; *)
(* find "bar" myt;; *)
(* find "zap" myt;; *)



(* let myt = add "z" 42 empty;; *)
(* let myt = add "y" 47 myt;; *)
(* let myt = add "x" 99 myt;; *)
(* let myt = add "w" 77 myt;; *)
(* let myt = add "v" 77 myt;; *)
(* let myt = add "u" 77 myt;; *)
(* let myt = add "t" 77 myt;; *)
(* let myt = add "s" 77 myt;; *)


(* let myt = add "z" 42 empty;; *)
(* let myt = add "x" 99 myt;; *)
(* let myt = add "y" 47 myt;; *)


(* let myt = add "a" 42 empty;; *)
(* let myt = add "b" 99 myt;; *)
(* let myt = add "c" 47 myt;; *)

(* let myt = add "a" 42 empty;; *)
(* let myt = add "c" 99 myt;; *)
(* let myt = add "b" 47 myt;; *)