Wojciech Jaworski / ENIAM

open Xstd

(* type 'a nodes = 'a IntMap.t *)
type 'a edges = 'a IntMap.t IntMap.t
type 'a tree = Tree of int * 'a * 'a tree list | Visited of int * 'a


(* Klucze zewnętrznego słownika określamy jako parent,
   klucze wewnętrznego słownika określamy jako child. *)

let empty = IntMap.empty

let add_new r parent child v =
  let children = try IntMap.find r parent with Not_found -> IntMap.empty in
  if IntMap.mem children child then
    failwith ("add_new: " ^ string_of_int parent ^ "-->" ^ string_of_int child ^ " is already in graph")
  else
    let children = IntMap.add children child v in
    IntMap.add r parent children

let add r parent child v =
  let children = try IntMap.find r parent with Not_found -> IntMap.empty in
  let children = IntMap.add children child v in
  IntMap.add r parent children

let add_inc r parent child v merge_fun =
  let children = try IntMap.find r parent with Not_found -> IntMap.empty in
  let children = IntMap.add_inc children child v (merge_fun v) in
  IntMap.add r parent children

let mem_parent r parent =
  IntMap.mem r parent

let rec find_descendants_rec r descendants parent =
  if IntSet.mem descendants parent then descendants else
  let descendants = IntSet.add descendants parent in
  let children = try IntMap.find r parent with Not_found -> IntMap.empty in
  IntMap.fold children descendants (fun descendants child _ -> find_descendants_rec r descendants child)

let find_descendants r parent =
  IntSet.remove (find_descendants_rec r IntSet.empty parent) parent

let test_reverse ra rb =
  IntMap.iter ra (fun parent_a children_a ->
    IntMap.iter children_a (fun parent_b v ->
      try
        let children_b = IntMap.find rb parent_b in
        if IntMap.mem children_b parent_a then ()
        else print_endline ("test_reverse_rel a: " ^ string_of_int parent_a)
      with Not_found -> print_endline ("test_reverse_rel b: " ^ string_of_int parent_b)))

let reverse r =
  IntMap.fold r IntMap.empty (fun r parent children ->
    IntMap.fold children r (fun r child v ->
      add r child parent v))

let sum ra rb merge_fun =
  IntMap.fold rb ra (fun r parent children ->
    IntMap.fold children r (fun r child v ->
      add_inc r parent child v merge_fun))

let select_childless selected r =
  IntMap.fold r selected (fun selected parent _ ->
    IntSet.remove selected parent)

let rec descendants_tree_rec r visited parent v =
  if IntSet.mem visited parent then visited, Visited(parent,v) else
  let visited = IntSet.add visited parent in
  let children = try IntMap.find r parent with Not_found -> IntMap.empty in
  let visited, l = IntMap.fold children (visited, []) (fun (visited,l) child v ->
    let visited,tree = descendants_tree_rec r visited child v in
    visited, tree :: l) in
  visited, Tree(parent,v,l)

let descendants_tree r parent v =
  snd (descendants_tree_rec r IntSet.empty parent v)

let rec create_spaces n =
  if n = 0 then "" else " " ^ (create_spaces (n - 1))

let rec print_tree_rec file node_fun level = function
    Tree(parent,cost,children) ->
      Printf.fprintf file "%s%s\n" (create_spaces (2*level)) (node_fun parent cost);
      Xlist.iter children (print_tree_rec file node_fun (level+1))
  | Visited(parent,cost) ->
      Printf.fprintf file "%sVISITED %s\n" (create_spaces (2*level)) (node_fun parent cost)

let print_tree file tree node_fun =
  print_tree_rec file node_fun 0 tree

let rec create_tree_xml file node_fun = function
    Tree(parent,cost,children) ->
      Xml.Element("node",node_fun parent cost,Xlist.map children (create_tree_xml file node_fun))
  | Visited(parent,cost) ->
      Xml.Element("node",node_fun parent cost,[])

let print_tree_xml file tree node_fun =
  let xml = create_tree_xml file node_fun tree in
  Printf.fprintf file "%s" (Xml.to_string_fmt xml)

let rec print_tree_as_graph_rec file node_fun edge_fun grand_parent = function
    Tree(parent,cost,children) ->
      Printf.fprintf file "  %d [label=\"%s\"]\n" parent (node_fun parent);
      if grand_parent <> (-1) then Printf.fprintf file "  %d -> %d [label=\"%s\"]\n" grand_parent parent (edge_fun cost);
      Xlist.iter children (print_tree_as_graph_rec file node_fun edge_fun parent)
  | Visited(parent,cost) ->
      Printf.fprintf file "  %d [label=\"%s\"]\n" parent (node_fun parent);
      if grand_parent <> (-1) then Printf.fprintf file "  %d -> %d [label=\"%s\"]\n" grand_parent parent (edge_fun cost)


let print_tree_as_graph path name tree node_fun edge_fun =
  File.file_out (path ^ name ^ ".gv") (fun file ->
    (*if lr then Printf.fprintf file "digraph G {\n  node [shape=box]\n  rankdir = LR\n"
    else*) Printf.fprintf file "digraph G {\n  node [shape=box]\n";
    print_tree_as_graph_rec file node_fun edge_fun (-1) tree;
    Printf.fprintf file "}\n");
  Sys.chdir path;
  ignore (Sys.command ("dot -Tpng " ^ name ^ ".gv -o " ^ name ^ ".png"));
  Sys.chdir "..";
  ()

let print file r cost_fun =
  IntMap.iter r (fun parent children ->
    IntMap.iter children (fun child cost ->
      Printf.fprintf file "%d\t%d\t%s\n" parent child (cost_fun cost)))

let rec find_connected_components_rec r is_included_fun visited conn parent =
  (* if parent = "28358" then print_endline "find_connected_components_rec 1"; *)
  if IntSet.mem visited parent || not (is_included_fun parent) then conn,visited else (
  (* if parent = "28358" then print_endline "find_connected_components_rec 2"; *)
  let conn = IntSet.add conn parent in
  let visited = IntSet.add visited parent in
  let l = try IntMap.find r parent with Not_found -> [] in
  (* if parent = "28358" then Printf.printf "find_connected_components_rec 3: |conn|=%d\n%!" (IntSet.size conn); *)
  Xlist.fold l (conn,visited) (fun (conn,visited) parent ->
    find_connected_components_rec r is_included_fun visited conn parent))

let find_connected_components r is_included_fun =
  let l,_ = IntMap.fold r ([],IntSet.empty) (fun (l,visited) parent _ ->
    if IntSet.mem visited parent then l,visited else
    let conn,visited = find_connected_components_rec r is_included_fun visited IntSet.empty parent in
    (* if IntSet.mem conn "28358" then print_endline "find_connected_components 1"; *)
    conn :: l, visited) in
  l

(* find_connected_components r threshold (fun id -> (IntMap.find synmap id).syn_no_hipo >= threshold) *)

(*
(* UWAGA: to działa gdy selektor jest monotoniczny względem relacji *)
let select_node_set edges selected =
  IntMap.fold edges IntMap.empty (fun edges pid nodes ->
    let nodes = IntSet.fold nodes IntSet.empty (fun nodes id ->
      if IntSet.mem selected id then IntSet.add nodes id else nodes) in
    if IntSet.is_empty nodes then edges else
    IntMap.add edges pid nodes)

let get_all_ids edges =
  IntMap.fold edges IntSet.empty (fun ids id nodes ->
    IntSet.fold nodes (IntSet.add ids id) IntSet.add)

let select_maximal_ids edges =
  let selected = get_all_ids edges in
  IntMap.fold edges selected (fun selected _ nodes ->
    IntSet.fold nodes selected (fun selected id ->
      IntSet.remove selected id))

let select_maximal_ids2 selected edges =
  IntMap.fold edges selected (fun selected _ nodes ->
    IntSet.fold nodes selected (fun selected id ->
      IntSet.remove selected id))

let print_graph path name lr edges node_fun =
  let ids = get_all_ids edges in
  let max_ids = select_maximal_ids edges in
  File.file_out (path ^ name ^ ".gv") (fun file ->
    if lr then Printf.fprintf file "digraph G {\n  node [shape=box]\n  rankdir = LR\n"
    else Printf.fprintf file "digraph G {\n  node [shape=box]\n";
    IntSet.iter ids (fun id ->
      if IntSet.mem max_ids id then
        Printf.fprintf file "  %d [color=\".7 .3 1.0\",style=filled,label=\"%s\"]\n" id (node_fun id)
      else Printf.fprintf file "  %d [label=\"%s\"]\n" id (node_fun id));
    IntMap.iter edges (fun id1 nodes ->
      IntSet.iter nodes (fun id2 ->
        Printf.fprintf file "  %d -> %d\n" id1 id2));
    Printf.fprintf file "}\n");
  Sys.chdir path;
  ignore (Sys.command ("dot -Tpng " ^ name ^ ".gv -o " ^ name ^ ".png"));
  Sys.chdir "..";
  ()
*)