type 'a tree = Node of 'a * 'a tree list
type 'a positioned_tree = PNode of 'a * float * 'a positioned_tree list
type extent = (float * float) list
type 'a abs_tree = ANode of 'a * float * 'a abs_tree list

let move_tree dx (PNode (label, x, children)) =
  PNode (label, x +. dx, children)

let move_extent dx : extent -> extent =
  List.map (fun (l, r) -> (l +. dx, r +. dx))

let rec merge_extents (e1 : extent) (e2 : extent) : extent =
  match e1, e2 with
  | [], qs -> qs
  | ps, [] -> ps
  | (l1, r1) :: ps, (l2, r2) :: qs ->
    (min l1 l2, max r1 r2) :: merge_extents ps qs

let merge_extent_list : extent list -> extent =
  List.fold_left merge_extents []

let fit (e1 : extent) (e2 : extent) : float =
  let rec go e1 e2 =
    match e1, e2 with
    | (_, r) :: ps, (l, _) :: qs -> max (r -. l +. 1.0) (go ps qs)
    | _, _ -> 0.0
  in
  go e1 e2

let fit_list_left (es : extent list) : float list =
  let rec go acc = function
    | [] -> []
    | e :: es ->
      let x = fit acc e in
      x :: go (merge_extents acc (move_extent x e)) es
  in
  go [] es

let fit_list (es : extent list) : float list =
  let lefts = fit_list_left es in
  let rights = List.rev (fit_list_left (List.rev_map (List.map (fun (l, r) -> (-.r, -.l))) es)) in
  let rights = List.map (fun x -> -.x) rights in
  List.map2 (fun l r -> (l +. r) /. 2.0) lefts rights

let rec design (Node (label, children) : 'a tree) : 'a positioned_tree * extent =
  let designed_children = List.map design children in
  let trees = List.map fst designed_children in
  let extents = List.map snd designed_children in
  let positions = fit_list extents in
  let positioned_trees = List.map2 move_tree positions trees in
  let positioned_extents = List.map2 move_extent positions extents in
  let result_extent = (0.0, 0.0) :: merge_extent_list positioned_extents in
  (PNode (label, 0.0, positioned_trees), result_extent)

let rec absolutize x (PNode (label, dx, children)) : 'a abs_tree =
  let x' = x +. dx in
  ANode (label, x', List.map (absolutize x') children)

let flatten tree =
  let rec go depth (ANode (label, x, children)) =
    (label, x, depth) :: List.concat_map (go (depth + 1)) children
  in
  go 0 tree

let edges tree =
  let rec go depth (ANode (_, x, children)) =
    let child_edges = List.map (fun (ANode (_, cx, _)) -> (x, depth, cx, depth + 1)) children in
    child_edges @ List.concat_map (go (depth + 1)) children
  in
  go 0 tree

module ReingoldTilford = struct
  type 'a rt_tree = {
    mutable x : float;
    mutable modifier : float;
    label : 'a;
    children : 'a rt_tree list;
    mutable thread : 'a rt_tree option;
    mutable ancestor : 'a rt_tree;
  }

  let rec make_rt tree =
    let Node (label, children) = tree in
    let rt_children = List.map make_rt children in
    let rec node = {
      x = 0.0; modifier = 0.0; label; children = rt_children;
      thread = None; ancestor = node;
    } in
    node

  let sibling_sep = 1.0
  let subtree_sep = 1.0

  let rec left_contour node depth acc =
    let acc = (depth, node) :: acc in
    match node.children with
    | [] -> (match node.thread with Some t -> left_contour t (depth + 1) acc | None -> acc)
    | first :: _ -> left_contour first (depth + 1) acc

  let rec right_contour node depth acc =
    let acc = (depth, node) :: acc in
    match List.rev node.children with
    | [] -> (match node.thread with Some t -> right_contour t (depth + 1) acc | None -> acc)
    | last :: _ -> right_contour last (depth + 1) acc

  let contour_offset left_tree right_tree =
    let left_c = left_contour left_tree 0 [] |> List.rev in
    let right_c = right_contour right_tree 0 [] |> List.rev in
    let rec compute lc rc l_off r_off max_sep =
      match lc, rc with
      | [], _ | _, [] -> max_sep
      | (ld, ln) :: lrest, (rd, rn) :: rrest when ld = rd ->
        let l_pos = ln.x +. l_off in
        let r_pos = rn.x +. r_off in
        let sep = l_pos -. r_pos +. sibling_sep in
        let l_off = l_off +. ln.modifier in
        let r_off = r_off +. rn.modifier in
        compute lrest rrest l_off r_off (max max_sep sep)
      | (ld, _) :: _, (rd, _) :: rrest when rd < ld ->
        compute lc rrest l_off r_off max_sep
      | (_, _) :: lrest, (_, _) :: _ ->
        compute lrest rc l_off r_off max_sep
    in
    compute right_c left_c 0.0 0.0 0.0

  let set_thread left_tree right_tree =
    let left_c = left_contour left_tree 0 [] in
    let right_c = right_contour right_tree 0 [] in
    let l_depth = match left_c with (d, _) :: _ -> d | [] -> -1 in
    let r_depth = match right_c with (d, _) :: _ -> d | [] -> -1 in
    if l_depth > r_depth then
      match right_c with (_, rnode) :: _ ->
        (match left_c |> List.find_opt (fun (d, _) -> d = r_depth + 1) with
         | Some (_, lnode) -> rnode.thread <- Some lnode
         | None -> ())
      | [] -> ()
    else if r_depth > l_depth then
      match left_c with (_, lnode) :: _ ->
        (match right_c |> List.find_opt (fun (d, _) -> d = l_depth + 1) with
         | Some (_, rnode) -> lnode.thread <- Some rnode
         | None -> ())
      | [] -> ()

  let rec first_walk node =
    match node.children with
    | [] -> ()
    | children ->
      List.iter first_walk children;
      let place_subtrees positioned current =
        match positioned with
        | [] -> [current]
        | prev :: _ ->
          let offset = contour_offset prev current in
          current.x <- prev.x +. offset +. subtree_sep;
          current.modifier <- current.x;
          set_thread prev current;
          current :: positioned
      in
      let _ = List.fold_left place_subtrees [] children in
      let first_child = List.hd children in
      let last_child = List.hd (List.rev children) in
      let mid = (first_child.x +. last_child.x) /. 2.0 in
      List.iter (fun c -> c.x <- c.x -. mid) children;
      List.iter (fun c -> c.modifier <- c.modifier -. mid) children

  let rec second_walk node offset =
    node.x <- node.x +. offset;
    List.iter (fun c -> second_walk c (offset +. node.modifier)) node.children

  let rec to_abs_tree node : 'a abs_tree =
    ANode (node.label, node.x, List.map to_abs_tree node.children)

  let layout tree =
    let rt = make_rt tree in
    first_walk rt;
    second_walk rt 0.0;
    to_abs_tree rt
end

let design_rt tree =
  ReingoldTilford.layout tree