(**************************************************************************)
(*                                                                        *)
(*  This file is part of OcamlGraph.                                      *)
(*                                                                        *)
(*  Copyright (C) 2009-2010                                               *)
(*    CEA (Commissariat � l'�nergie Atomique)                             *)
(*                                                                        *)
(*  you can redistribute it and/or modify it under the terms of the GNU   *)
(*  Lesser General Public License as published by the Free Software       *)
(*  Foundation, version 2.1, with a linking exception.                    *)
(*                                                                        *)
(*  It is distributed in the hope that it will be useful,                 *)
(*  but WITHOUT ANY WARRANTY; without even the implied warranty of        *)
(*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *)
(*  GNU Lesser General Public License for more details.                   *)
(*                                                                        *)
(*  See the file ../LICENSE for more details.                             *)
(*                                                                        *)
(*  Authors:                                                              *)
(*    - Julien Signoles  (Julien.Signoles@cea.fr)                         *)
(*    - Jean-Denis Koeck (jdkoeck@gmail.com)                              *)
(*    - Benoit Bataille  (benoit.bataille@gmail.com)                      *)
(*                                                                        *)
(**************************************************************************)

module type S = sig

  module Tree: Graphviz.GraphWithDotAttrs

  module TreeManipulation : sig
    type t
    val get_structure : t -> Tree.t
    val get_tree_vertices : Tree.V.label -> t -> Tree.V.t list
    val get_graph_vertex : Tree.V.t -> t -> Tree.V.label
    val is_ghost_node : Tree.V.t -> t -> bool
    val is_ghost_edge : Tree.E.t -> t -> bool
  end

  type cluster = string
  type graph_layout

  class tree_model :
    graph_layout ->
    TreeManipulation.t ->
    [Tree.V.t, Tree.E.t, cluster] DGraphModel.abstract_model

  val tree : unit -> TreeManipulation.t

end

module Build
    (G: Sig.G)
    (T: Graphviz.GraphWithDotAttrs with type V.label = G.V.t)
    (TM: DGraphSubTree.S with type Tree.t = T.t
                          and type Tree.V.t = T.V.t
                          and type Tree.E.t = T.E.t) =
struct

  module TreeManipulation = TM
  type cluster = string

  module X = XDot.Make(T)
  type graph_layout = X.graph_layout

  class tree_model layout tree
    : [ T.V.t, T.E.t, cluster ] DGraphModel.abstract_model
    =
    let tree_structure = TM.get_structure tree in
    object

      (* Iterators *)
      method iter_edges f =
        T.iter_edges
          (fun v1 v2 ->
             if not (TM.is_ghost_node v1 tree && TM.is_ghost_node v2 tree) then
               f v1 v2)
          tree_structure

      method iter_edges_e f =
        T.iter_edges_e
          (fun e -> if not (TM.is_ghost_edge e tree) then f e)
          tree_structure

      method iter_pred f v =
        T.iter_pred
          (fun v -> if not (TM.is_ghost_node v tree) then f v)
          tree_structure v

      method iter_pred_e f v =
        T.iter_pred_e
          (fun e -> if not (TM.is_ghost_edge e tree) then f e)
          tree_structure v

      method iter_succ f =
        T.iter_succ
          (fun v -> if not (TM.is_ghost_node v tree) then f v)
          tree_structure

      method iter_succ_e f =
        T.iter_succ_e
          (fun e -> if not (TM.is_ghost_edge e tree) then f e)
          tree_structure

      method iter_vertex f =
        T.iter_vertex
          (fun v -> if not (TM.is_ghost_node v tree) then f v)
          tree_structure

      method iter_associated_vertex f v =
        let origin_vertex = TM.get_graph_vertex v tree in
        List.iter
          (fun v -> if not (TM.is_ghost_node v tree) then f v)
          (TM.get_tree_vertices origin_vertex tree)

      method iter_clusters f =
        Hashtbl.iter (fun k _ -> f k) layout.X.cluster_layouts

      (* Membership functions *)
      method find_edge =
        try T.find_edge tree_structure
        with Not_found -> assert false

      method mem_edge = T.mem_edge tree_structure
      method mem_edge_e = T.mem_edge_e tree_structure
      method mem_vertex = T.mem_vertex tree_structure
      method src = T.E.src
      method dst = T.E.dst

      (* Layout *)
      method bounding_box = layout.X.bbox

      method get_vertex_layout v =
        try X.HV.find layout.X.vertex_layouts v
        with Not_found -> assert false

      method get_edge_layout e =
        try X.HE.find e layout.X.edge_layouts
        with Not_found -> assert false

      method get_cluster_layout c =
        try Hashtbl.find layout.X.cluster_layouts c
        with Not_found -> assert false

    end

end

module SubTreeMake(G: Graphviz.GraphWithDotAttrs) = struct

  module T = Imperative.Digraph.Abstract(G.V)
  module TM = DGraphSubTree.Make(G)(T)

  let tree_ref : TM.t option ref = ref None
  let tree () = match !tree_ref with None -> assert false | Some t -> t

  let graph_ref: G.t option ref = ref None
  let graph () = match !graph_ref with None -> assert false | Some g -> g

  module Tree = struct

    include T

    let graph_attributes _ = G.graph_attributes (graph ())
    let default_vertex_attributes _ = G.default_vertex_attributes (graph ())
    let default_edge_attributes _ = G.default_edge_attributes (graph ())

    let cpt = ref 0
    let name_table = Hashtbl.create 97
    let vertex_name v =
      try Hashtbl.find name_table v
      with Not_found ->
        incr cpt;
        Hashtbl.add name_table v (string_of_int !cpt);
        string_of_int !cpt

    let vertex_attributes v =
      let t = tree () in
      if TM.is_ghost_node v t then [ `Style `Invis ]
      else G.vertex_attributes (TM.get_graph_vertex v t)

    let edge_attributes e =
      let t = tree () in
      if TM.is_ghost_node (T.E.src e) t || TM.is_ghost_node (T.E.dst e) t then
        [ `Style `Dashed; `Dir `None ]
      else
        G.edge_attributes
          (G.find_edge
             (graph ())
             (TM.get_graph_vertex (T.E.src e) t)
             (TM.get_graph_vertex (T.E.dst e) t))

    let get_subgraph v =
      let t = tree () in
      if TM.is_ghost_node v t then None
      else G.get_subgraph (TM.get_graph_vertex v t)

  end

  include Build(G)(Tree)(TM)

  module TreeLayout =
    DGraphTreeLayout.Make
      (Tree)
      (struct let is_ghost_node v = TM.is_ghost_node v (tree ()) end)

  let from_graph
      ?(depth_forward=2)
      ?(depth_backward=2)
      ~fontMeasure
      g
      v
    =
    (* Generate subtree *)
    let t = TM.make g v depth_forward depth_backward in
    tree_ref := Some t;
    graph_ref := Some g;
    let layout =
      TreeLayout.from_tree ~fontMeasure (TM.get_structure t) (TM.get_root t)
    in
    new tree_model layout t

end

module SubTreeDotModelMake = struct

  module T = Imperative.Digraph.Abstract(DGraphModel.DotG.V)
  module TM = DGraphSubTree.Make_from_dot_model(T)

  let tree_ref : TM.t option ref = ref None
  let tree () = match !tree_ref with None -> assert false | Some t -> t

  module TreeLayout =
    DGraphTreeLayout.MakeFromDotModel
      (T)
      (struct let is_ghost_node v = TM.is_ghost_node v (tree ()) end)

  include TreeLayout
  include Build(DGraphModel.DotG)(Tree)(TM)

  let from_model ?(depth_forward=2) ?(depth_backward=2) model v =
    let t = TM.make model v depth_forward depth_backward in
    tree_ref := Some t;
    let tree_structure = TM.get_structure t in
    let layout = from_model tree_structure (TM.get_root t) model in
    new tree_model layout t

end
