open Graphics ;;

type 'a tree = Empty | Leaf of 'a | Node of 'a * 'a tree * 'a tree ;;

(* 
STRUCT : (digit, xcoord, ycoord) 
*)

let rec pw x n = match n with
    | 0 -> 1
    | 1 -> x
    | k when k mod 2 = 0 -> let res = pw x (n/2) in res*res
    | k -> let res = pw x (n/2) in res*res*x ;;

let rec depth_of_tree t = match t with
    | Leaf _ -> 1
    | Node (_, g, d) -> 1 + max (depth_of_tree g) (depth_of_tree d) 
    | Empty -> 0;;

let fill_data te ystep sx sy r =
    let depth = depth_of_tree te in
    let res = Array.make (depth+1) [] in
    let rec aux t cur_x cur_d spacing pcx pcy = match t with
        | Node (x, g, d) -> begin
            aux g (cur_x - spacing) (cur_d+1) (spacing/2) cur_x (sy - r - 20 - ystep * cur_d);
            res.(cur_d) <- (((x, (pcx, pcy)), (cur_x, sy - r - 20 - ystep * cur_d)))::(res.(cur_d));
            aux d (cur_x + spacing) (cur_d+1) (spacing/2) cur_x (sy - r - 20 - ystep * cur_d);
        end
        | Leaf x -> begin
            res.(cur_d) <- (((x, (pcx, pcy)), (cur_x, sy - r - ystep * cur_d)))::(res.(cur_d));
        end
        | Empty -> ()
    in aux te (sx/2) 0 (r/2 + r * ((pw 2 (depth-1)) - 1)) (-1) (-1); res ;;

let rec ln10 n = match n with
    | k when k < 0 -> failwith "Are you sure about that ?"
    | k when k < 10 -> 0
    | k -> 1 + ln10 (k/10) ;;

let delta i j =
    if i = j then 1 else 0 ;;


let draw_integer x0 y n0 r =
  (* 7-seg display *)
  let n = ref n0 in
  let size = ln10 n0 in
  let len = r/3 in
  let offset = size*(len*11/7)/2 in
  for i = 0 to size do
      let x = x0 + offset - i*(len*11/7) in
      if Array.mem (!n mod 10) [|0; 4; 5; 6; 7; 8; 9|] then
          draw_poly_line [|(x-len/2, y+len); (x-len/2, y)|];

      if Array.mem (!n mod 10) [|0; 2; 3; 5; 6; 7; 8; 9|] then
          draw_poly_line [|(x-len/2, y+len); (x+len/2, y+len)|];

      if Array.mem (!n mod 10) [|0; 1; 2; 3; 4; 7; 8; 9|] then
          draw_poly_line [|(x+len/2, y+len); (x+len/2, y)|];

      if Array.mem (!n mod 10) [|2; 3; 4; 5; 6; 8; 9|] then
          draw_poly_line [|(x-len/2, y); (x+len/2, y)|];

      if Array.mem (!n mod 10) [|0; 1; 3; 4; 5; 6; 7; 8; 9|] then
          draw_poly_line [|(x+len/2, y-len); (x+len/2, y)|];

      if Array.mem (!n mod 10) [|0; 2; 3; 5; 6; 8; 9|] then
          draw_poly_line [|(x-len/2, y-len); (x+len/2, y-len)|];

      if Array.mem (!n mod 10) [|0; 2; 6; 8|] then
          draw_poly_line [|(x-len/2, y-len); (x-len/2, y)|];

      n := !n/10;
  done ;;

let rec draw_list l d r = match l with
    | [] -> ()
    | h::t -> begin
        set_color (rgb 192 192 192);
        fill_circle (fst (snd h)) (snd (snd h)) r;
        set_color black;
        draw_circle (fst (snd h)) (snd (snd h)) r;
        moveto (fst (snd h)) (snd (snd h)); 
        set_color (rgb 32 192 32);
        draw_integer (fst (snd h)) (snd (snd h)) (fst (fst h)) r; 
        draw_list t d r 
    end;;

let connect l0 =
    let rec aux l = match l with
        | [] -> ()
        | ((_, (xf, yf)), (x, y))::t -> 
            if xf >= 0 && yf >= 0 then begin
                set_color (rgb 192 192 192);
                draw_poly_line [|(xf, yf); (x, y)|];
                aux t
            end
    in aux l0 ;;

let even_more_pretty_printing t r ystep skip =

    let sx = Graphics.size_x () in
    let sy = Graphics.size_y () in

    let graphdata = fill_data t ystep sx sy (6*r/10) in
    (* graphdata is a ((int * (int * int)) * (int * int)) list array *)
    (* <==> ((value, (parent_x, parent_y)), (this_x, this_y)) *)
    if skip = false then begin
        set_color (rgb 192 192 192);
        set_line_width 15 ;
        for dpth = 1 to (Array.length graphdata -1) do
            connect graphdata.(dpth-1);
        done;

        set_line_width (max 1 (r/6)) ;
        for dpth = 0 to (Array.length graphdata -1) do
            draw_list graphdata.(dpth) dpth r
        done;

        (*
        let halt = ref false in 
        while !halt = false do
            Unix.sleepf 0.1 ;
            Unix.sleepf 2.0 ;
            halt := true;
        done;*)
    end;
    graphdata ;;


let generate_full_tree d = 
    let rec aux n = match n with
    | 0 -> Leaf (Random.int 1000)
    | k -> begin 
        Node (Random.int 1000, aux (n-1), aux (n-1))
    end 
    in aux d ;;

let generate_some_tree maxd nodechance leafchance = 
    let rec aux n = match n with
    | 0 -> if (Random.int 101 < leafchance) then Leaf (Random.int 100) else Empty
    | k when k = maxd -> Node (Random.int 1000, aux (maxd-1), aux (maxd-1))
    | k -> begin 
        match Random.int 101 with
        | k when k <= nodechance -> Node (Random.int 1000, aux (n-1), aux (n-1))
        | k -> if (Random.int 101 < leafchance) then Leaf (Random.int 1000) else Empty
    end 
    in aux maxd ;;

let rec nth l n = match l with
    | [] -> failwith "Out of range"
    | h::t when n = 0 -> h
    | h::t -> nth t (n-1) ;;

let even_more_fancy_dfs_prefixe t graphdata r tts rfound gfound bfound rmark gmark bmark =
    let d = depth_of_tree t in
    let count_per_depth = Array.make d 0 in
    let rec aux tr dpth = 
        match tr with
        | Empty -> ()
        | Leaf _ -> begin
            let data = nth graphdata.(dpth) (List.length graphdata.(dpth) - count_per_depth.(dpth) - 1) in
            count_per_depth.(dpth) <- count_per_depth.(dpth) + 1;
            set_color (rgb rfound gfound bfound);
            draw_circle (fst (snd data)) (snd (snd data)) r;
            Unix.sleepf tts;
            set_color (rgb rmark gmark bmark);
            draw_circle (fst (snd data)) (snd (snd data)) r;
        end
        | Node (_, g, d) -> begin
            let data = nth graphdata.(dpth) (List.length graphdata.(dpth) - count_per_depth.(dpth) - 1) in
            count_per_depth.(dpth) <- count_per_depth.(dpth) + 1;
            set_color (rgb rfound gfound bfound);
            draw_circle (fst (snd data)) (snd (snd data)) r;
            Unix.sleepf tts;
            set_color (rgb rmark gmark bmark);
            draw_circle (fst (snd data)) (snd (snd data)) r;

            aux g (dpth+1);
            aux d (dpth+1);
        end
    in aux t 0 ;;


let even_more_fancy_dfs_infixe t graphdata r tts rfound gfound bfound rmark gmark bmark =
    let rec aux tr dpth os = 
        match tr with
        | Empty -> ()
        | Leaf _ -> begin
            let data = nth graphdata.(dpth) (List.length graphdata.(dpth) - os - 1) in
            set_color (rgb rfound gfound bfound);
            draw_circle (fst (snd data)) (snd (snd data)) r;
            Unix.sleepf tts;
            set_color (rgb rmark gmark bmark);
            draw_circle (fst (snd data)) (snd (snd data)) r;
        end
        | Node (_, g, d) -> begin
            aux g (dpth+1) (2*os);

            let data = nth graphdata.(dpth) (List.length graphdata.(dpth) - os - 1) in
            set_color (rgb rfound gfound bfound);
            draw_circle (fst (snd data)) (snd (snd data)) r;
            Unix.sleepf tts;
            set_color (rgb rmark gmark bmark);
            draw_circle (fst (snd data)) (snd (snd data)) r;

            aux d (dpth+1) (2*os + 1);
        end
    in aux t 0 0 ;;


let even_more_fancy_dfs_postfixe t graphdata r tts rfound gfound bfound rmark gmark bmark =
    let rec aux tr dpth os = 
        match tr with
        | Empty -> ()
        | Leaf _ -> begin
            let data = nth graphdata.(dpth) (List.length graphdata.(dpth) - os - 1) in
            set_color (rgb rfound gfound bfound);
            draw_circle (fst (snd data)) (snd (snd data)) r;
            Unix.sleepf tts;
            set_color (rgb rmark gmark bmark);
            draw_circle (fst (snd data)) (snd (snd data)) r;
        end
        | Node (_, g, d) -> begin
            aux g (dpth+1) (2*os);
            aux d (dpth+1) (2*os + 1);

            let data = nth graphdata.(dpth) (List.length graphdata.(dpth) - os - 1) in
            set_color (rgb rfound gfound bfound);
            draw_circle (fst (snd data)) (snd (snd data)) r;
            Unix.sleepf tts;
            set_color (rgb rmark gmark bmark);
            draw_circle (fst (snd data)) (snd (snd data)) r;
        end
    in aux t 0 0 ;;

(* NEW VERSION *)

type pt = {x : int ; y :int} ;;
type node_data = {tag : int ; parent : pt ; self : pt} ;;

type 'a data_tree = Nothing | Tail of node_data | Cross of node_data * 'a data_tree * 'a data_tree ;;
(* changing names to avoid confusion *)

let count_per_floor tr =
    let d = depth_of_tree tr in
    let res = Array.make d 0 in
    let rec aux tr dpth = match tr with
        | Empty -> ()
        | Leaf _ -> res.(dpth) <- res.(dpth) + 1
        | Node (_, g, d) -> res.(dpth) <- res.(dpth) + 1 ; aux g (dpth+1) ; aux d (dpth+1)
    in aux tr 0 ; res ;;

let showtree tdt r = 
    let rec aux t side = match t with
    | Nothing -> ()
    | Tail data -> begin
        set_line_width 9;
        if side = 1 then set_color (rgb 200 48 48) else set_color (rgb 48 48 200) ;
        draw_poly_line [|(data.parent.x, data.parent.y); (data.self.x, data.self.y)|];

        set_color (rgb 192 192 192);
        fill_circle data.self.x data.self.y r;

        set_color (rgb 32 192 32);
        set_line_width 7;
        draw_circle data.self.x data.self.y r;

        set_color black;
        set_line_width 5;
        draw_integer data.self.x data.self.y data.tag r;
    end 
    | Cross (data, g, d) -> begin
        set_line_width 9;
        if side = 1 then set_color (rgb 200 48 48) else set_color (rgb 48 48 200) ;
        draw_poly_line [|(data.parent.x, data.parent.y); (data.self.x, data.self.y)|];

        aux g (-1);
        aux d 1;

        set_color (rgb 192 192 192);
        fill_circle data.self.x data.self.y r;

        set_color (rgb 192 192 32);
        set_line_width 7;
        draw_circle data.self.x data.self.y r;

        set_color black;
        set_line_width 5;
        draw_integer data.self.x data.self.y data.tag r;
    end 
    in aux tdt 0 ;;

let coords_on_segment a b divsize k =
    if divsize <> 0 then
        a + k*(b-a)/divsize 
    else (a + b)/2 ;;

let max_of_arr a = 
    let m = ref a.(0) in
    for i = 1 to (Array.length a -1) do
        if !m < a.(i) then m := a.(i)
    done; !m ;;

let pretty_tree_printing_new_version tr r ystep win_w win_h display =
    let d = depth_of_tree tr in
    let amt_per_floor = count_per_floor tr in

    let visited_fl = Array.make d 0 in
    (* visited.(x) count the number of already visited nodes in floor x *)

    let rec build_data_tree tr dpth parent_xy = 
        match tr with
        | Empty -> Nothing
        | Leaf x -> begin
            (*let self = {x = win_w/2 - (14*r/6)*amt_per_floor.(dpth)/2 + (14*r/6)*visited_fl.(dpth); y = win_h - r/2 - (dpth)*ystep} in*)
            let self = {x = coords_on_segment (max r (win_w/2 - 2*r*(pw 2 dpth))) (min (win_w - r) (win_w/2 + 2*r*(pw 2 dpth))) (amt_per_floor.(dpth)-1) visited_fl.(dpth); y = win_h - r/2 - (dpth)*ystep} in
            visited_fl.(dpth) <- visited_fl.(dpth) + 1;
            let data = {tag = x ; parent = parent_xy ; self = self} in
            Tail (data)
        end
        | Node (x, g, d) -> begin
            (*let self = {x = win_w/2 - (14*r/6)*amt_per_floor.(dpth)/2 + (14*r/6)*visited_fl.(dpth); y = win_h - r/2 - (dpth)*ystep} in*)
            let self = {x = coords_on_segment (max r (win_w/2 - 2*r*(pw 2 dpth))) (min (win_w - r) (win_w/2 + 2*r*(pw 2 dpth))) (amt_per_floor.(dpth)-1) visited_fl.(dpth); y = win_h - r/2 - (dpth)*ystep} in
            visited_fl.(dpth) <- visited_fl.(dpth) + 1;
            if dpth <> 0 then begin
                let data = {tag = x ; parent = parent_xy ; self = self} in
                let arg_left = build_data_tree g (dpth+1) self in
                let arg_right = build_data_tree d (dpth+1) self in
                (* PS : this is a good example of OCaml evaluating its arguments from right to left *)
                (* if the recursive call were to be directly inside the constructor, the displayed tree would be reversed *)
                Cross (data, arg_left, arg_right)
            end else begin
                let data = {tag = x ; parent = self ; self = self} in
                let arg_left = build_data_tree g (dpth+1) self in
                let arg_right = build_data_tree d (dpth+1) self in
                Cross (data, arg_left, arg_right)
            end
        end
    in 
    let treedata = build_data_tree tr 0 {x = win_w/2 ; y = win_h - r} in 
    if display then showtree treedata r; treedata ;;

(* ABR things *)

let identity n = n ;;

let rec insert_abr tr e = match tr with
    | Empty -> Node (e, Empty, Empty)
    | Leaf t when e < t -> Node (t, (Node (e, Empty, Empty)), Empty)
    | Leaf t -> Node (t, Empty, (Node (e, Empty, Empty)))
    | Node (x, g, d) when e < x -> Node (x, insert_abr g e, d) 
    | Node (x, g, d) -> Node (x, g, insert_abr d e) ;;

let successive_insert () =
    let cur_tree = ref (Empty) in
    open_graph " 1400x1000" ;
    set_window_title "Trees" ;
    try
        while true do
            Stdlib.print_endline "What element would you like to insert ? (crash to terminate)";
            let elt = Scanf.bscanf Scanf.Scanning.stdin "%d\n" identity in

            open_graph " 1200x1000" ;
            set_window_title "Trees" ;

            cur_tree := insert_abr !cur_tree elt;
            (*ignore (pretty_tree_printing_new_version !cur_tree 40 100 1200 1000 true)*)
            ignore (even_more_pretty_printing !cur_tree 20 100 false);
        done;
        ()
    with
        | Stdlib.Scanf.Scan_failure _ -> ignore (even_more_pretty_printing !cur_tree 20 100 false) ;close_graph () ;;

(* --------------------------------------| TESTS |-------------------------------------- *)
Random.self_init () ;;
(*
open_graph " 1800x1000" ;;
set_window_title "Trees" ;;
ignore (pretty_tree_printing_new_version (Node (0, Node (1, (Node (0, Node (1, Empty, Empty), Node (2, Empty, Empty))), Empty), Node (2, Empty, (Node (0, Node (1, Empty, Empty), Node (2, Empty, Empty)))))) 40 150 1800 1000 true) ;;
ignore (Scanf.bscanf Scanf.Scanning.stdin "%d\n" identity) ;;
close_graph () ;;
failwith "E" ;;
*)
successive_insert () ;;

open_graph " 1800x1000" ;;
set_window_title "Trees" ;;

let tt = generate_some_tree 5 75 100 ;;

ignore (pretty_tree_printing_new_version tt 40 150 1800 1000 true) ;;

ignore (Scanf.bscanf Scanf.Scanning.stdin "%d\n" identity) ;;
(*
let gdata = even_more_pretty_printing tt 30 150 false ;;

even_more_fancy_dfs_prefixe tt gdata 30 0.2 255 255 32 32 32 255 ;;*)

close_graph () ;;

(* compilation command : ocamlfind ocamlc -linkpkg -package unix -linkpkg -package graphics trees.ml *)
print_int 0 ;;
print_char '\n' ;;