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 5 ;
        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 = match t with
    | Nothing -> ()
    | Tail data -> begin
        set_line_width 9;
        set_color (rgb 48 48 48);
        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;
        set_color (rgb 48 48 48);
        draw_poly_line [|(data.parent.x, data.parent.y); (data.self.x, data.self.y)|];

        aux g;
        aux d;

        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 ;;

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
                Cross (data, build_data_tree g (dpth+1) self, build_data_tree d (dpth+1) self)
            end else begin
                let data = {tag = x ; parent = self ; self = self} in
                Cross (data, build_data_tree g (dpth+1) self, build_data_tree d (dpth+1) self)
            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 ;;

(* --------------------------------------| TESTS |-------------------------------------- *)
Random.self_init () ;;

let identity n = n ;;

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' ;;