Ax/lib/draw_tree.ml

open Base
open Math

type dir = [`H | `V]
type fig = Rect of Dim.t * Color.t
type path = int list Lazy.t ref

type tree = (arm, leaf) Generic_tree.t
and arm  = [`H of path | `V of path | `Uninit]
and leaf = path * Event.Id.t * fig
and t = tree

let create_path () =
  ref (lazy (failwith "Path.to_list: path is uninitialized!"))

let uninit : t =
  Arm(`Uninit, [])
let arm dir ts : t =
  Arm((match dir with
      | `H -> `H(create_path ())
      | `V -> `V(create_path ())), ts)
let figure fg : t =
  Leaf((create_path (), Event.Id.gen (), fg))

let is_init : t -> bool = function
  | Arm(`Uninit, _) -> false
  | _ -> true

let rec event_ids : t -> _ = function
  | Arm(_, ts) ->
     Sequence.(bind (of_list ts)) ~f:event_ids
  | Leaf((_, ev_id, _)) ->
     Sequence.singleton ev_id

module Zipper = Generic_tree.Zipper
type zipper = (arm, leaf) Zipper.t

let fig_to_string = function
  | Rect(dim, color) ->
     Printf.sprintf "Rect(%s,%s)"
       (Dim.to_string dim)
       (Color.to_string color)

let fig_equal fg1 fg2 = match fg1, fg2 with
  | Rect(d1, c1), Rect(d2, c2) ->
     Dim.equal d1 d2 && Color.equal c1 c2

include Generic_tree.Make_show_eq
    (struct
      type t = arm
      let to_string = function
        | `H _ -> "H"
        | `V _ -> "V"
        | `Uninit -> "Uninit"
      let equal a b = match a, b with
        | `H _,    `H _    -> true
        | `V _,    `V _    -> true
        | `Uninit, `Uninit -> true
        | _,       _       -> false
    end)
    (struct
      type t = leaf
      let to_string (_, _, fg) = fig_to_string fg
      let equal (_, _, fg1) (_, _, fg2) = fig_equal fg1 fg2
    end)


(*** paths ***)

module Path = struct
  type t = path
  let root () = ref (Lazy.from_val [])
  let of_list l = ref (Lazy.from_val l)
  let of_list_rev l = ref (lazy (List.rev l))
  let to_list p = Lazy.force !p
  let set_rev p l = p := lazy (List.rev l)
  let to_string p =
    Printf.sprintf "path[%s]" @@
    String.concat ~sep:";" @@
    List.map (to_list p) ~f:Int.to_string
  let equal a b =
    Polymorphic_compare.(to_list a = to_list b)
end

let path : t -> path = function
  | Leaf((path, _, _)) -> path
  | Arm((`H(path) | `V(path)), _) -> path
  | Arm(`Uninit, _) -> raise (Invalid_argument("Draw_tree.path called on uninit"))

let rec refresh_paths ?(rev_prefix=[]) (t : t) =
  let () = Path.set_rev (path t) rev_prefix in
  match t with
  | Leaf _ -> ()
  | Arm(_, ts) ->
     List.iteri ts
       ~f:(fun i t ->
         refresh_paths ~rev_prefix:(i::rev_prefix) t)


(*** updating ***)

module Update = struct
  type t = Set of fig | Mount of tree | Unmount

  let to_string = function
    | Set(fig) -> Printf.sprintf "Set(%s)" (fig_to_string fig)
    | Mount(t) -> Printf.sprintf "Mount(%s)" (to_string t)
    | Unmount  -> "Unmount"

  let equal a b = match a, b with
    | Set(fg1),  Set(fg2)  -> fig_equal fg1 fg2
    | Mount(t1), Mount(t2) -> equal t1 t2
    | Unmount,   Unmount   -> true
    | _,         _         -> false

  let fail_message = function
    | Mount _ -> "cannot mount already-initialized tree"
    | Set _   -> "cannot update non-figure tree"
    | Unmount -> "cannot umount already-uninitialized tree"

  let apply_exn ~path (upd : t) : zipper -> zipper =
    Zipper.modify ~f:(fun tree ->
        match upd, tree with
        | Set(new_fig), Leaf((path, id, _)) ->
           Leaf((path, id, new_fig))
        | Mount(tree), Arm(`Uninit, _) ->
           let () = refresh_paths ~rev_prefix:(List.rev path) tree in
           tree
        | Unmount, _ when is_init tree ->
           uninit
        | _, _ ->
           failwith (fail_message upd))
end

type updates = (Path.t * Update.t) Sequence.t

let apply_updates_exn upds t : t =
  Zipper.to_tree @@ snd @@
  Sequence.fold upds ~init:([], Zipper.of_tree t)
    ~f:(fun (p, zip) (p', upd) ->
        let p' = Path.to_list p' in
        (p', (zip
              |> Zipper.nav_exn ~prev:p ~targ:p'
              |> Update.apply_exn ~path:p' upd)))

(*** generic geometry calculation ***)

module type GEOMETRY_MONOID = sig
  type t
  val concat : t Sequence.t -> t
  val rect : AABB.t -> Color.t -> Event.Id.t -> t
end

let geometry
      (type a)
      (module A : GEOMETRY_MONOID with type t = a)
      (tree : t) : a =

  let rec compute : t -> (int * int * (AABB.t -> a)) = function
    | Leaf(_, ev_id, Rect({w;h} as dim, fill)) ->
       let go bounds = A.rect (AABB.center bounds dim) fill ev_id in
       (w, h, go)

    | Arm(`Uninit, _) ->
       (0, 0, fun _ -> A.concat Sequence.empty)

    | Arm((`H _ | `V _) as ty, ts) ->
       let axis1 x y = match ty with `H _ -> x | `V _ -> y in
       let axis2 x y = match ty with `H _ -> y | `V _ -> x in
       let pair a1 a2 = match ty with `H _ -> (a1, a2) | `V _ -> (a2, a1) in

       let (sizes_funcs, max_size, total_size) =
         List.fold_right ts ~init:([], 0, 0)
           ~f:(fun t (s_f, max_sz, tot_sz) ->
               let (w, h, func) = compute t in
               ((axis1 w h, func)::s_f,
                max (axis2 w h) max_sz,
                axis1 w h + tot_sz))
       in

       let go {bx;by;bh;bw} =
         let init = (axis1 bw bh - total_size) / 2 in
         Sequence.folding_map (Sequence.of_list sizes_funcs) ~init
           ~f:(fun off (size, func) ->
               (off + size,
                func { bx = bx + axis1 off 0
                     ; by = by + axis2 off 0
                     ; bw = axis1 size bw
                     ; bh = axis2 size bh }))
         |> A.concat
       in

       let (w, h) = pair total_size max_size in
       (w, h, go)
  in
  let (bw, bh, func) = compute tree in
  func {bx=0;by=0;bw;bh}


(*** rendering ***)

module Draw = struct
  type t = Rect of AABB.t * Color.t

  let equal x y = match x, y with
    | Rect(b1, c1), Rect(b2, c2) ->
       AABB.equal b1 b2 && Color.equal c1 c2

  let to_string = function
    | Rect(b, c) -> Printf.sprintf "Rect(%s, %s)"
                      (AABB.to_string b)
                      (Color.to_string c)
end

module Draw_monoid = struct
  type nonrec t = Draw.t Sequence.t
  let concat = Sequence.concat
  let rect aabb color _ = Sequence.singleton (Draw.Rect(aabb, color))
end

let render tree : Draw.t Sequence.t =
  geometry (module Draw_monoid) tree


(*** collision ***)

module Collide_monoid = struct
  type t = Pos.t -> Event.Id.t option
  let concat ts pos = Sequence.find_map ts ~f:(fun f -> f pos)
  let rect bounds _ id pos = Option.some_if (AABB.contains bounds pos) id
end

let capture_point ~pos tree : Event.Id.t option =
  geometry (module Collide_monoid) tree pos