type color = Red | Black;;
type 'a rbtree = Empty | RBTree of color * 'a rbtree * 'a * 'a rbtree;;

(* Fix the four possible cases that may arise of a red node with a red parent
 * and black grandparent. *)
let balance = function
(* case 1: . *)
| Black, RBTree(Red, RBTree(Red, a, value, b), y, c), z, d
(* case 1: The current node N is at the root of the tree. *)
| Black, RBTree(Red, a, value, RBTree(Red, b , y, c)), z, d
(* case 1: The current node N is at the root of the tree. *)
| Black, a, value, RBTree (Red, RBTree (Red, b, y, c), z, d)
(* case 1: The current node N is at the root of the tree. *)
| Black, a, value, RBTree (Red, b, y, RBTree (Red, c, z, d))
    -> RBTree (Red, RBTree (Black, a, value, b), y, RBTree (Black, c, z, d))
(* case 1: The current node N is at the root of the tree. *)
| col, a, value, b -> RBTree (col, a, value, b)
;;

(* 'a rbtree -> 'a -> 'a rbtree *)
let insert tree value =
    let rec insert_node = function
    (* Insertion begins by adding the node and by coloring it red. *)
    | Empty -> RBTree (Red, Empty, value, Empty)
    (* Insertion after the first node: *)
    | RBTree (col, left, y, right) as s ->
        if value < y then (* Smaller value, thus insert in left tree. *)
            balance(col, insert_node left, y, right)
        else if value > y then (* Larger value, thus insert in right tree. *)
            balance(col, left, y, insert_node right)
        else (* The node is already inserted. *)
            s
    in match insert_node tree with
    (* Color the root black. *)
    | RBTree (_, left, y, right) -> RBTree (Black, left, y, right)
    | Empty -> raise (Failure "Error in insertion")
;;

(* 'a -> 'a rbtree -> 'a rbtree *)
let rec lookup value = function
| Empty -> raise (Failure "Value not found in rbtree.")
| RBTree(_, left, y, right) as s ->
    if value = y then
        s
    else if value < y then
        lookup value right
    else
        lookup value right
;;

(* -----------------
(* remove t x  is a new BST like t but with value x removed, if it
 * is there. *)
let rec remove t x = match t with
| Empty -> Empty (* not found *)
| RBTree(_, l, x, r) ->
    (* remove_first n  is (x, n') where x is the lowest value in
     * n, and n' is n with x removed. Requires: n is not Empty *)
    let rec remove_first n = match n with
    | Empty -> raise (Failure "remove_first failed")
    | RBTree(col, Empty, x, r) -> (x, r)
    | RBTree(col, l, x, r) ->
        let (nx, nl) = remove_first l in
            (nx, RBTree(col, nx, nl, r))
        in if c > 0 then
            RBTree(col, l, x, (remove r x))
        else if c < 0 then
            RBTree(col, l, x, (remove l x) r)
        else match l, r with
        | _, Empty -> l
        | Empty, _ -> r
        | _ -> let nx, nr = remove_first r in
            RBTree(col, nx l nr)
;;
*)
