open PMap;;
open List;;

exception Cykliczne

(* tworzy graf z danych wejściowych, z założeniem że dany wierzcholek
może występować wiele razy w liscie (np [(1,[2]); (1,[3]);...] *)

let make_graph lt =
    match lt with
    | [] -> (PMap.empty, PMap.empty)
    (*zwracamy pusty graf*)
    | lt ->
      let graph = PMap.empty and colors = PMap.empty in
      let helper (graph, colors) (el, rel) =
        if PMap.mem el graph then (*przypadek, że wierzcholek już występuje *)
          let old = PMap.find el graph in
          let graph = PMap.remove el graph in
          let graph = PMap.add el (old @ rel) graph in
          (graph, colors)
        else
          let graph = PMap.add el rel graph
          and colors = PMap.add el 0 colors in
          (graph, colors)
      in List.fold_left helper (graph, colors) lt


(*sortowanie topologiczne wykorzystujące DFS *)

let top_sort (graph,colors) =
  let rec dfs el (graph, colors) res =
    if not (PMap.mem el colors) then
      let colors = PMap.add el 2 colors in
      ((graph, colors), el::res)
      (* wierzcholek nie jest w grafie colors wiec nie wychodza z niego zadne
      krawedzie i ustawiamy jako odwiedzony*)
    else
      let col = PMap.find el colors in
      if col = 1 then raise Cykliczne (*wierzcołek z którego wyszliśmy *)
      else if col = 2 then ((graph, colors), res) (* wierzchołek w którym byliśmy *)
      else  (*wierzcholek nieodwiedzony*)
        let colors = PMap.add el 1 colors in
        let helper ((graph,colors), res) el =
        dfs el (graph,colors) res in
        let ((graph,colors), res) = List.fold_left helper ((graph,colors), res) (PMap.find el graph) in
        let colors = PMap.add el 2 colors in
        ((graph,colors), el::res)
  in
  let ((graph,colors), res) = PMap.foldi (fun el _ ((graph,colors), res) -> dfs el (graph,colors) res) graph ((graph, colors), [])  in
  res

let topol lt =
  let (graph,colors) = make_graph lt in
  top_sort (graph, colors)