;; open Assert

(* A type of integer-carrying trees *)
type tree = 
  | Empty
  | Node of tree * int * tree

(* counts the number of nodes in the tree *)
let rec size (t:tree) : int =
  begin match t with
    | Empty -> 0
    | Node(l,_,r) -> 1 + (size l) + (size r)
  end

(* counts the longest path from the root to a leaf *)
let rec height (t:tree) : int =
  begin match t with
    | Empty -> 0
    | Node(l,_,r) -> 1 + max (height l) (height r)
  end

(* returns the preorder traversal of the tree *)
let rec preorder (t:tree) : int list =
  begin match t with
    | Empty -> []
    | Node(l,n,r) -> n::(preorder l)@(preorder r)
  end

(* returns the in order traversal of the tree *)
let rec inorder (t:tree) : int list =
  begin match t with
    | Empty -> []
    | Node(l,n,r) -> (inorder l)@(n::(inorder r))
  end

(* returns the post-order traversal of the tree *)
let rec postorder (t:tree) : int list =
  begin match t with
    | Empty -> []
    | Node(l,n,r) -> (postorder l)@(postorder r)@[n]
  end

(***********************)
(* tree pretty printer *)
(***********************)

(* NOTE:
   This portion of the code uses OCaml features that we haven't yet
   covered in class.  We include it here so that you can play around
   with trees and print them out in a nice way.

   The function string_of_tree pretty prints a tree to a string.
   All of the rest of the functions and datastructures are 
   auxilliary to this main function.

   Note that string_of_tree assumes that the integers at the nodes 
   of the tree are in the rang [0-9] so that they print as a 
   single character.
*)

(* a representation of pictures *)
type picture = string array

(* a picture with some extra information:
    rootx is the x coordinate (starting from 0)
    of the root character of this tree 
*)
type tree_block = {
  width : int;
  height : int;
  rootx : int;
  picture : picture;
}

let char_of_int (i:int) : char =
  Char.chr ((Char.code '0') + i)

(* creates a new tree block with a root value v 
   rooted at rx *)
let new_tb (w:int) (h:int) (rx:int) (v:int) =
  let pic = Array.init h (fun _ -> String.make w ' ') in
  let _ = String.set (pic.(0)) rx  (char_of_int v) in
    {width = w;
     height = h;
     rootx = rx;
     picture = pic;}

(* inserts newlines into a treeblock to serialize it
   as a single string *)
let string_of_tree_block (b:tree_block) : string =
  String.concat "\n" (Array.to_list (b.picture))

(* an exception that is raised if there is overlap 
   detected when trying to superimpose two trees *)
exception Overlap

(* copies one tree_block into another; doesn't check for overlap 
   x = starting x position
   dy = amount to shift the copied block vertically
*)
let overwrite (dst:tree_block) (src:tree_block) (x:int) (dy:int) : unit =
  for i = 0 to src.height - 1 do
    String.blit (src.picture.(i)) x (dst.picture.(i+dy)) x src.width
  done

(* copies on tree_block into another, checking for collisions (overlap)
   dx and dy are offsets to shift the source block relative to the target *)
let try_to_write (dst:tree_block) (src:tree_block) (dx:int) (dy:int) : unit =
  for i = 0 to src.height - 1 do
    for j = 0 to src.width - 1 do
      let y = i+dy in
      let x = j+dx in
      let src_char = String.get (src.picture.(i)) j in
	if    ((String.get (dst.picture.(y)) x) <> ' ') 
	   && (src_char <> ' ')
	then
	  raise Overlap
	else
	  String.set (dst.picture.(y)) x src_char
    done
  done

(* draws an left-slanting edge of length el starting at root
   rx *)
let draw_left_edge (dst:tree_block) (el:int) (rx:int) : unit =
  for i = 1 to el do
    String.set dst.picture.(i) (rx - i) '/'
  done

(* draws a right-slanting edge of length el starting at root
   rx *)
let draw_right_edge (dst:tree_block) (el:int) (rx:int) : unit =
  for i = 1 to el do
    String.set dst.picture.(i) (rx + i) '\\'
  done

(* superimposes two tree_blocks and gives them a new root r.
   it tries to connect the two trees to the root using an edge
   of length el.  Raises 'Overlap' if there is a collision in 
   this process *)
let superimpose (b1:tree_block) (b2:tree_block) (r:int) (el:int) : tree_block =
  let dx = el * 2 + 2 in
  let dy = el + 1 in 
  let new_height = dy + (max b1.height b2.height) in
  let (fixed_block, shifted_block, shift_amt, new_width, new_rootx) = 
    if (b1.rootx + dx) >= b2.rootx then
      (* b1 is not shifted *)
      let amt = b1.rootx + dx - b2.rootx in
      let w = max b1.width (b1.rootx + dx +(b2.width - b2.rootx)) in
      let rx = b1.rootx + dy in
	(b1, b2, amt, w, rx)
    else
      (* b1 is shifted *)
      let amt = b2.rootx - (b1.rootx + dx) in
      let w = max (b1.width + amt) b2.width in
      let rx = b1.rootx + amt + dy in
	(b2, b1, amt, w, rx)
  in
  let new_block = new_tb new_width new_height new_rootx r in
  let _ = overwrite new_block fixed_block 0 dy in
  let _ = try_to_write new_block shifted_block shift_amt dy in
  let _ = draw_left_edge new_block el new_rootx in
  let _ = draw_right_edge new_block el new_rootx in
    new_block

(* merges two tree blocks by trying to connect them to a new
   root.  It starts with the smallest possible edge length
   and increases it until no collision overlaps are detected *)
let merge (b1:tree_block) (b2:tree_block) (r:int) : tree_block =
  let rec loop (el:int) : tree_block =
    try superimpose b1 b2 r el with
      | Overlap -> loop (el+1)
  in
    loop 1


(* makes a tree_block into the left child of a new root r *)
let make_left_child (t:tree_block) (r:int) : tree_block =
  let new_height = t.height + 2 in
  let new_width = max (t.width) (t.rootx + 3) in
  let new_rootx = t.rootx + 2 in
  let new_block = new_tb new_width new_height new_rootx r in
  let _ = overwrite new_block t 0 2 in
  let _ = draw_left_edge new_block 1 new_rootx in
    new_block

(* makes a tree_block into the right child of a new root r *)
let make_right_child (t:tree_block) (r:int) : tree_block =
  let new_height = t.height + 2 in
  let shift_amt = max 0 (2 - t.rootx) in
  let new_width = t.width + shift_amt in
  let new_rootx = max 0 (t.rootx - 2) in
  let new_block = new_tb new_width new_height new_rootx r in
  let _ = try_to_write new_block t shift_amt 2 in
  let _ = draw_right_edge new_block 1 new_rootx in
    new_block

(* creates an empty tree_block *)
let empty_block () : tree_block =
  {height=0; width=0; rootx=0; picture=Array.make 0 ""}

(* creates an tree_block with just one node, which is 
   both the root and leaf *)
let leaf_block (r:int) : tree_block =
  {height=1; width=1; rootx=0; 
   picture=Array.make 1 (String.make 1 (char_of_int r));}

(* recursively constructs a tree_block from a tree
   using 'merge' to find an optimal non-overlapping 
   picture for the children *)
let rec block_of_tree (t:tree) : tree_block =
  begin match t with 
    | Empty -> empty_block ()
    | Node(Empty, r, Empty) -> leaf_block r
    | Node(Empty, r, rt) -> 
	let rtb = block_of_tree rt in
	  make_right_child rtb r
    | Node(lt, r, Empty) ->
	let ltb = block_of_tree lt in
	  make_left_child ltb r
    | Node(lt, r, rt) ->
	let ltb = block_of_tree lt in
	let rtb = block_of_tree rt in
	  merge ltb rtb r
  end

(* The main function:  
   assumes that the nodes are in [0-9]
*)
let string_of_tree (t:tree) : string =
  string_of_tree_block (block_of_tree t)
	
	
(************** Example Trees *********************)	
	

(* helper function for constructing trees *)
let leaf (i:int) : tree =
  Node(Empty, i, Empty)

(* Some example trees *)
let t1 = 
  Node(Node(leaf 0,
	    2,
	    leaf 1),
       3,
       Node(leaf 3,
	    2,
	    leaf 1))


let t2 = Node(t1, 5, t1)

let t3 = Node(Empty, 7,
	      Node(
		Node(
		  Node(
		    Node(
		      Empty,
		      3,
		      Empty),
		    4,
		    Empty),
		  5,
		  Empty),
		6,
		Empty))


let t4 = Node(Node(Empty,
		   5,
		   Node(Empty,
			6, 
			Node(Empty,
			     8,
			     leaf 9)))
		, 7,
	      Empty)

let t5 = 
  Node(Node(leaf 2, 1, Empty),
       0,
       Node(leaf 7, 6, leaf 8))
			
			
			
			

(*  Example tree from lecture 7 *)				
										    
let lec07_tree =
  Node(Node(leaf 0, 1, leaf 3),
       5,
       Node(Empty, 7,
	    Node(leaf 8, 9, Empty)))

;; print_endline "********** lec07_tree *************"
;; print_endline (string_of_tree lec07_tree)

(* an pre-order search through an arbitrary tree.
   returns true if the tree contains n 
 *)
let rec contains (t:tree) (n:int) : bool =
  begin match t with
    | Empty -> false
    | Node(lt, x, rt) -> 
	x = n || (contains lt n) || (contains rt n)
  end

let test () : bool = contains lec07_tree 9
;; run_test "contains 9" test
let test () : bool = not (contains lec07_tree 12) 
;; run_test "contains 12" test


(***********************)
(* BINARY SEARCH TREES *)
(***********************)

(* Determines whether t contains n *)
let rec lookup (t:tree) (n:int) : bool =
  begin match t with
    | Empty -> false
    | Node(lt, x, rt) ->
			if n = x then true else 
				if n < x then lookup lt n
				else lookup rt n
  end

let test () : bool = lookup lec07_tree 9
;; run_test "lookup 9" test
let test () : bool = not (lookup lec07_tree 12) 
;; run_test "lookup 12" test







(* helper functions for writing is_bst *)
(* (tree_less t n) is true when all nodes of t are 
   strictly less than n
*)
let rec tree_less (t:tree) (n:int) : bool =
  begin match t with
    | Empty -> true
    | Node(lt, x, rt) -> 
	    x < n && (tree_less lt n) && (tree_less rt n)
  end

(* (tree_gtr t n) is true when all nodes of t are 
   strictly greater than n 
*)
let rec tree_gtr (t:tree) (n:int) : bool =
  begin match t with
    | Empty -> true
    | Node(lt, x, rt) ->
	    x > n && (tree_gtr lt n) && (tree_gtr rt n)
  end

(* determines whether t satisfies the binary *)
(* search tree invariant *)
let rec is_bst (t:tree) : bool =
  begin match t with
    | Empty -> true
    | Node(lt, x, rt) ->
	    is_bst lt && is_bst rt &&
	    (tree_less lt x) && (tree_gtr rt x)
  end
	
let test () : bool = 
	is_bst lec07_tree
;; run_test "is_bst tree" test		
	
	
	
	

(* Inserts n into the binary search tree t *)
let rec insert (t:tree) (n:int) : tree =
  begin match t with
    | Empty -> Node (Empty, n, Empty)
    | Node(lt, x, rt) ->
			if n < x then 
				 Node(insert lt n, x, rt) 
			else if n > x then 
				 Node(lt, x, insert rt n)
			else t	
  end
	
let insert_result =
  Node(Node(leaf 0, 1, Node (Empty, 3, leaf 4)),
       5,
       Node(Empty, 7,
	    Node(leaf 8, 9, Empty)))	
			
let test () : bool = 
	insert lec07_tree 4 = insert_result
;; run_test "insert 4"	test		
let test () : bool = 
	is_bst (insert lec07_tree 4) 
;; run_test "insert is BST" test	
		
	
	
	
	
	
	

(* returns the maximum integer in a *NONEMPTY* binary
   search tree t  *)
let rec tree_max (t:tree) : int =
  begin match t with
    | Empty -> failwith "tree_max called on empty tree"
    | Node(_,x,Empty) -> x
    | Node(_,_,rt) -> tree_max rt
  end
	
let test () : bool = 
	tree_max lec07_tree = 9
;; run_test "tree_max t" test		







(* returns a binary search tree that has the same set of 
   nodes as t except with n removed (if it's there)
*)
let rec delete (t:tree) (n:int) : tree =
  begin match t with
    | Empty -> Empty
    | Node(lt,x,rt) ->
	if x = n then
	  begin match (lt, rt) with
			| (Empty, Empty) -> Empty
			| (Empty, rt) -> rt 
			| (lt, Empty) -> lt
			| (lt, rt) -> let y = tree_max lt in
			   (Node (delete lt y, y, rt))
		end
	else
	  if n < x then Node(delete lt n, x, rt)
	  else Node(lt, x, delete rt n)
  end

let delete_three_result = 
  Node(Node(leaf 0, 1, Empty),
       5,
       Node(Empty, 7,
	    Node(leaf 8, 9, Empty)))
let test () : bool = 
	delete lec07_tree 3 = delete_three_result
;; run_test "delete 3 (no children)" test

let delete_seven_result = 
  Node(Node(leaf 0, 1, leaf 3),
       5,
	    Node(leaf 8, 9, Empty))
			
let test () : bool = 
	delete lec07_tree 7 = delete_seven_result
;; run_test "delete 7 (one child)" test					

let delete_five_result = 
  Node(Node(leaf 0, 1, Empty),
       3,
       Node(Empty, 7,
	    Node(leaf 8, 9, Empty)))
let test () : bool = 
	delete lec07_tree 5 = delete_five_result
;; run_test "delete 5 (two children)" test


let random_bst =
  let _ = Random.self_init () in
  let rec loop (t:tree) (n:int) : tree =
    if n=0 then t else 
      let x = Random.int 10000000 in
	loop (insert t x) (n-1)
  in
    loop Empty