open Graph
open Printf
open Tools

module Int_set = Set.Make(struct
    type t = int
    let compare = compare
  end)

module Int_map = Map.Make(struct
    type t = int
    let compare = compare
  end)

type 'a network = {
  graph: 'a graph;
  origin: id;
  destination: id;
}

type flow = {
  current: int;
  capacity: int;
}

type path = (id * id) list

(* Converts an int graph to a flow graph with the current capacity at 0 for each arc *)
(* int graph -> flow graph *)
let initialize_graph gr = gmap gr (fun c -> { current = 0; capacity = c })

(* adds an arc to the graph only if the value is different from 0 *)
(* int graph -> int -> int -> int -> int graph *)
let add_res_arc gr origin destination value = match value with
  | 0 -> gr
  | _ -> add_arc gr origin destination value

(* Builds a residual graph from a flow graph *)
(* flow graph -> int graph *)
let build_res_graph gr = let gr_res = clone_nodes gr in
  e_fold gr (fun g origin destination fl ->
      let forward = fl.capacity - fl.current in
      let backward = fl.current in
      let new_g = add_res_arc g origin destination forward in
      add_res_arc new_g destination origin backward )
    gr_res

(* Converts the predecessors Map to a path from origin to destination. An empty list is returned when no path is found *)
(* Int_map -> id -> id -> path *)
let predecessors_to_path predecessors origin destination =
  let rec loop o d acu = if o = d then
      acu
    else
      match Int_map.find_opt d predecessors with
      | None -> []
      | Some x -> let new_acu = (x, d) :: acu in
        if x = o then new_acu else
          loop o x new_acu
  in
  loop origin destination []

(* finds a path from the origin to the destination in the network graph *)
(* int network -> path *)
let find_path res_network = let queue = Queue.create () in
  let () = Queue.add res_network.origin queue in
  let visited = Int_set.add res_network.origin Int_set.empty in
  let rec find_path_rec q v pred = match Queue.take_opt q with
    | None -> predecessors_to_path pred res_network.origin res_network.destination
    | Some x -> if x = res_network.destination then
        predecessors_to_path pred res_network.origin res_network.destination
      else
        let arcs = out_arcs res_network.graph x in
        let filtered_arcs = List.filter (fun (dest, _) -> not (Int_set.mem dest v) ) arcs in
        let successors_seq = List.to_seq ( 
            List.map (fun (dest, _) -> dest) filtered_arcs
          ) in
        let predecessors_seq = List.to_seq ( 
            List.map (fun (dest, _) -> (dest, x)) filtered_arcs
          ) in 
        let new_pred = Int_map.add_seq predecessors_seq pred in
        let new_v = Int_set.add_seq successors_seq v in
        let () = Queue.add_seq q successors_seq in
        find_path_rec q new_v new_pred
  in 
  find_path_rec queue visited Int_map.empty

(* Finds the minimum capacity on the path in the given graph *)
(* int graph -> path -> int *)
let find_min gr path = List.fold_left (fun current_min (s, d) -> 
    let current = find_arc gr s d in
    match current with
    | None -> raise Not_found
    | Some x -> min x current_min)
    max_int path 

(* Apply the given augmentation or reduction along the path on the flow graph *)
(* flow graph -> path -> int -> flow graph *)
let apply_path gr path aug = List.fold_left (fun new_gr (s, d) -> 
    match find_arc new_gr s d with
    | None -> (match find_arc new_gr d s with
        | None -> raise Not_found
        | Some x -> new_arc new_gr d s {current = x.current - aug; capacity = x.capacity})
    | Some x -> new_arc new_gr s d {current = x.current + aug; capacity = x.capacity})
    gr path

(* Gets the max flow from the graph *)
(* flow network -> int *)
let get_max_flow flow_network = let arcs = out_arcs flow_network.graph flow_network.origin in
  List.fold_left (fun current (dest, flow) -> current + flow.current ) 0 arcs

(* Runs the ff algorithm on the given network *)
(* int network -> flow network *)
let run_ff network = let flow_graph = initialize_graph network.graph in
  let rec loop fl = 
    let res_graph = build_res_graph fl in
    let path = find_path {graph = res_graph; origin = network.origin; destination = network.destination } in
    match path with
    | [] -> {graph = fl; origin = network.origin; destination = network.destination }
    | _ ->
      let aug = find_min res_graph path in
      let new_fl = apply_path fl path aug in
      loop new_fl
  in loop flow_graph