include Graph_sig

module Make(L: Map.OrderedType) = struct
  module LabelMap = Map.Make(L)
  type label = L.t

  type 'v labels   = ('v * int) LabelMap.t
  type 'v vertices = 'v array
  type 'e edges    = 'e option array array

  type ('v, 'e) t = {labels: 'v labels; vertices: 'v vertices; edges: 'e edges}

  let make vertices edges =
    let labels =
      List.to_seq vertices
      |> Seq.mapi (fun i (l, v) -> l, (v, i))
      |> LabelMap.of_seq in
    let vertices = Array.of_list (List.map snd vertices) in
    let len = Array.length vertices in
    let edge arr (i, o, l) =
      let get x = snd (LabelMap.find x labels) in
      arr.(get i).(get o) <- Some l in
    let edges =
      let arr = Array.make_matrix len len None in
      List.iter (edge arr) edges;
      arr in
    {labels; vertices; edges}

  let make' vertices = make (List.map (fun x -> x, x) vertices)

  let make_e edges =
    let vertices = List.concat_map (function a, b, _ -> [a; b]) edges in
    make' vertices edges

  let count_vertices {vertices; _} = Array.length vertices

  let count_edges {edges; _} =
    let count = ref 0 in
    Array.iter (Array.iter (Option.iter (fun _ -> incr count))) edges;
    !count

  let rev ({edges; _} as graph) =
    let len = Array.length graph.vertices in
    let edges =
      Array.init len (fun j -> Array.init len (fun i -> edges.(i).(j))) in
    {graph with edges}


  let iter_dfs f {vertices; edges; labels} label =
    let seen = Array.(make (length vertices) false) in
    let rec go i =
      let edge o = function
        | None   -> ()
        | Some e ->
            let vi = vertices.(i) in
            let vo = vertices.(o) in
            go o;
            f vi vo e in
      if not seen.(i) then begin
        seen.(i) <- true;
        Array.iteri edge edges.(i)
      end in
    go (snd (LabelMap.find label labels))
end


let%test_module _ = (module struct
  open Make(String)
  let lmap lst = List.to_seq lst |> LabelMap.of_seq

  let print_graph_from = iter_dfs (Format.printf "%s -> %s: %d\n")

  let one_edge = make_e ["a", "b", 7]
  let triangle = make_e ["a", "b", 7; "b", "c", 10; "c", "a", 61]
  let vee      = make_e ["a", "b", 1; "a", "c", 2]

  let%test _ =
    make [] [] =
    {labels   = LabelMap.empty;
     vertices = [||];
     edges    = [||]}

  let%test _ = one_edge =
    {labels   = lmap ["a", ("a", 0); "b", ("b", 1)];
     vertices = [|"a"; "b"|];
     edges    = [|[|None; Some 7|]; [|None; None|]|]}

  let%expect_test _ =
    print_graph_from one_edge "a";
    [%expect{| a -> b: 7 |}]

  let%expect_test _ =
    print_graph_from triangle "b";
    [%expect{|
      a -> b: 7
      c -> a: 61
      b -> c: 10 |}]

  let%expect_test _ =
    print_graph_from (rev triangle) "b";
    [%expect{|
      c -> b: 10
      a -> c: 61
      b -> a: 7 |}]

  let%expect_test _ =
    print_graph_from vee "a";
    [%expect{|
      a -> b: 1
      a -> c: 2 |}]
end)