(* This file is free software, part of containers. See file "license" for more details. *)

(** {1 Lazy Tree Structure}
    This structure can be used to represent trees and directed
    graphs (as infinite trees) in a lazy fashion. It
    is a structural type. *)

type 'a iter = ('a -> unit) -> unit
type 'a gen = unit -> 'a option
type 'a printer = Format.formatter -> 'a -> unit
type +'a t = unit -> [ `Nil | `Node of 'a * 'a t list ]

let empty () = `Nil

let is_empty t =
  match t () with
  | `Nil -> true
  | `Node _ -> false

let singleton x () = `Node (x, [])
let node x l () = `Node (x, l)
let node1 x t () = `Node (x, [ t ])
let node2 x t1 t2 () = `Node (x, [ t1; t2 ])

let rec fold f acc t =
  match t () with
  | `Nil -> acc
  | `Node (x, l) ->
    let acc = f acc x in
    List.fold_left (fold f) acc l

let rec iter f t =
  match t () with
  | `Nil -> ()
  | `Node (x, l) ->
    f x;
    List.iter (iter f) l

let size t = fold (fun n _ -> n + 1) 0 t

let height t =
  let rec aux t k =
    match t () with
    | `Nil -> k 0
    | `Node (_, l) -> aux_l 0 l k
  and aux_l acc l k =
    match l with
    | [] -> k acc
    | t' :: l' -> aux t' (fun n -> aux_l (max acc n) l' k)
  in
  aux t (fun x -> x)

let rec map f t () =
  match t () with
  | `Nil -> `Nil
  | `Node (x, l) -> `Node (f x, List.map (map f) l)

let ( >|= ) t f = map f t

let rec cut_depth n t () =
  match t () with
  | `Nil -> `Nil
  | `Node _ when n = 0 -> `Nil
  | `Node (x, l) -> `Node (x, List.map (cut_depth (n - 1)) l)

(** {2 Graph Traversals} *)

(** Abstract Set structure *)
class type ['a] pset =
  object
    method add : 'a -> 'a pset
    method mem : 'a -> bool
  end

let set_of_cmp (type elt) ~cmp () =
  let module S = Set.Make (struct
    type t = elt

    let compare = cmp
  end) in
  object
    val s = S.empty
    method add x = {<s = S.add x s>}
    method mem x = S.mem x s
  end

let _nil () = Seq.Nil
let _cons x l = Seq.Cons (x, l)

let dfs ~pset t =
  let rec dfs pset stack () =
    match stack with
    | [] -> Seq.Nil
    | `Explore t :: stack' ->
      (match t () with
      | `Nil -> dfs pset stack' ()
      | `Node (x, _) when pset#mem x -> dfs pset stack' () (* loop *)
      | `Node (x, l) ->
        let pset' = pset#add x in
        let stack' =
          List.rev_append
            (List.rev_map (fun x -> `Explore x) l)
            (`Exit x :: stack')
        in
        _cons (`Enter x) (dfs pset' stack'))
    | `Exit x :: stack' -> _cons (`Exit x) (dfs pset stack')
  in
  dfs pset [ `Explore t ]

(** Functional queues for BFS *)
module FQ = struct
  type 'a t = { hd: 'a list; tl: 'a list }

  exception Empty

  (* invariant: if hd=[], then tl=[] *)
  let _make hd tl =
    match hd with
    | [] -> { hd = List.rev tl; tl = [] }
    | _ :: _ -> { hd; tl }

  let empty = _make [] []

  let list_is_empty = function
    | [] -> true
    | _ :: _ -> false

  let is_empty q = list_is_empty q.hd
  let push q x = _make q.hd (x :: q.tl)

  let pop_exn q =
    match q.hd with
    | [] ->
      assert (list_is_empty q.tl);
      raise Empty
    | x :: hd' ->
      let q' = _make hd' q.tl in
      x, q'
end

let bfs ~pset t =
  let rec bfs pset q () =
    if FQ.is_empty q then
      Seq.Nil
    else (
      let t, q' = FQ.pop_exn q in
      match t () with
      | `Nil -> bfs pset q' ()
      | `Node (x, _) when pset#mem x -> bfs pset q' () (* loop *)
      | `Node (x, l) ->
        let q' = List.fold_left FQ.push q' l in
        let pset' = pset#add x in
        _cons x (bfs pset' q')
    )
  in
  bfs pset (FQ.push FQ.empty t)

let rec force t : [ `Nil | `Node of 'a * 'b list ] as 'b =
  match t () with
  | `Nil -> `Nil
  | `Node (x, l) -> `Node (x, List.map force l)

let find ~pset f t =
  let rec _find_kl f l =
    match l () with
    | Seq.Nil -> None
    | Seq.Cons (x, l') ->
      (match f x with
      | None -> _find_kl f l'
      | Some _ as res -> res)
  in
  _find_kl f (bfs ~pset t)

(** {2 Pretty-printing} *)

let pp pp_x fmt t =
  (* at depth [lvl] *)
  let rec pp fmt t =
    match t with
    | `Nil -> ()
    | `Node (x, children) ->
      let children = filter children in
      (match children with
      | [] -> pp_x fmt x
      | _ :: _ ->
        Format.fprintf fmt "@[<v2>(@[<hov0>%a@]%a)@]" pp_x x pp_children
          children)
  and filter l =
    let l =
      List.fold_left
        (fun acc c ->
          match c () with
          | `Nil -> acc
          | `Node _ as sub -> sub :: acc)
        [] l
    in
    List.rev l
  and pp_children fmt children =
    (* remove empty children *)
    List.iter
      (fun c ->
        Format.fprintf fmt "@,";
        pp fmt c)
      children
  in
  pp fmt (t ());
  ()

(** {2 Pretty printing in the DOT (graphviz) format} *)

module Dot = struct
  type attribute =
    [ `Color of string
    | `Shape of string
    | `Weight of int
    | `Style of string
    | `Label of string
    | `Id of string
    | `Other of string * string ]
  (** Dot attributes for nodes *)

  type graph = string * attribute list t list
  (** A dot graph is a name, plus a list of trees labelled with attributes *)

  let mk_id format =
    let buf = Buffer.create 64 in
    Printf.kbprintf (fun _ -> `Id (Buffer.contents buf)) buf format

  let mk_label format =
    let buf = Buffer.create 64 in
    Printf.kbprintf (fun _ -> `Label (Buffer.contents buf)) buf format

  let make ~name l = name, l
  let singleton ~name t = name, [ t ]

  (* find and remove the `Id attribute, if any *)
  let rec _find_id acc l =
    match l with
    | [] -> raise Not_found
    | `Id n :: l' -> n, List.rev_append acc l'
    | x :: l' -> _find_id (x :: acc) l'

  let _pp_attr fmt attr =
    match attr with
    | `Color c -> Format.fprintf fmt "color=%s" c
    | `Shape s -> Format.fprintf fmt "shape=%s" s
    | `Weight w -> Format.fprintf fmt "weight=%d" w
    | `Style s -> Format.fprintf fmt "style=%s" s
    | `Label l -> Format.fprintf fmt "label=\"%s\"" l
    | `Other (name, value) -> Format.fprintf fmt "%s=\"%s\"" name value
    | `Id _ -> ()
  (* should not be here *)

  let rec _pp_attrs fmt l =
    match l with
    | [] -> ()
    | [ x ] -> _pp_attr fmt x
    | x :: l' ->
      _pp_attr fmt x;
      Format.pp_print_char fmt ',';
      _pp_attrs fmt l'

  let pp out (name, l) =
    (* nodes already printed *)
    let tbl = Hashtbl.create 32 in
    (* fresh name generator *)
    let new_name =
      let n = ref 0 in
      fun () ->
        let s = Printf.sprintf "node_%d" !n in
        incr n;
        s
    in
    (* the name for some node is either defined, either a fresh random
        name *)
    let get_name x = try _find_id [] x with Not_found -> new_name (), x in
    (* recursive printing (bfs) *)
    let rec aux q =
      if FQ.is_empty q then
        ()
      else (
        let (parent, x), q' = FQ.pop_exn q in
        let q' = pp_node q' ?parent x in
        aux q'
      )
    and pp_node q ?parent t =
      match t () with
      | `Nil -> q
      | `Node (x, l) ->
        let name, attrs = get_name x in
        (match parent with
        | None -> ()
        | Some n -> Format.fprintf out "  %s -> %s;@," n name);
        if not (Hashtbl.mem tbl name) then (
          Hashtbl.add tbl name ();
          Format.fprintf out "@[%s [%a];@]@," name _pp_attrs attrs;
          List.fold_left (fun q y -> FQ.push q (Some name, y)) q l
        ) else
          q
    in
    let q = List.fold_left (fun q y -> FQ.push q (None, y)) FQ.empty l in
    (* preamble *)
    Format.fprintf out "@[<hv 2>digraph \"%s\" {@," name;
    aux q;
    Format.fprintf out "}@]@.";
    ()

  let pp_single name out t = pp out (singleton ~name t)

  let print_to_file filename g =
    let oc = open_out filename in
    let fmt = Format.formatter_of_out_channel oc in
    try
      pp fmt g;
      Format.pp_print_flush fmt ();
      close_out oc
    with e ->
      close_out oc;
      raise e

  let to_file ?(name = "graph") filename trees =
    let g = make ~name trees in
    print_to_file filename g
end