defmodule Graph.Pathfinding do
  @moduledoc """
  This module contains implementation code for path finding algorithms used by `libgraph`.
  """
  import Graph.Utils, only: [vertex_id: 1, edge_weight: 3]

  @doc """
  Finds the shortest path between `a` and `b` as a list of vertices.
  Returns `nil` if no path can be found.

  The shortest path is calculated here by using a cost function to choose
  which path to explore next. The cost function in Dijkstra's algorithm is
  `weight(E(A, B))+lower_bound(E(A, B))` where `lower_bound(E(A, B))` is always 0.
  """
  def dijkstra(%Graph{} = g, a, b) do
    a_star(g, a, b, fn _v -> 0 end)
  end

  @doc """
  Finds the shortest path between `a` and `b` as a list of vertices.
  Returns `nil` if no path can be found.

  This implementation takes a heuristic function which allows you to
  calculate the lower bound cost of a given vertex `v`. The algorithm
  then uses that lower bound function to determine which path to explore
  next in the graph.

  The `dijkstra` function is simply `a_star` where the heuristic function
  always returns 0, and thus the next vertex is chosen based on the weight of
  the edge between it and the current vertex.
  """
  def a_star(%Graph{vertices: vs, out_edges: oe} = g, a, b, hfun) when is_function(hfun, 1) do
    with a_id  <- vertex_id(a),
         b_id  <- vertex_id(b),
         {:ok, a_out} <- Map.fetch(oe, a_id) do
      tree = Graph.new |> Graph.add_vertex(a_id)
      q = PriorityQueue.new
      q =
        a_out
        |> Stream.map(fn id -> {id, cost(g, a_id, id, hfun)} end)
        |> Enum.reduce(q, fn {id, cost}, q -> PriorityQueue.push(q, {a_id, id}, cost) end)
      case do_bfs(q, g, b_id, tree, hfun) do
        nil ->
          nil
        path ->
          for id <- path, do: Map.get(vs, id)
      end
    else
      _ -> nil
    end
  end

  @doc """
  Finds all paths between `a` and `b`, each path as a list of vertices.
  Returns `nil` if no path can be found.
  """
  def all(%Graph{vertices: vs, out_edges: oe} = g, a, b) do
    with a_id  <- vertex_id(a),
         b_id  <- vertex_id(b),
         {:ok, a_out} <- Map.fetch(oe, a_id) do
      case dfs(g, a_out, b_id, [a_id], []) do
        nil ->
          []
        paths ->
          paths
          |> Enum.map(fn path -> Enum.map(path, &Map.get(vs, &1)) end)
      end
    else
      _ -> []
    end
  end

  ## Private

  defp cost(%Graph{vertices: vs} = g, v1_id, v2_id, hfun) do
    edge_weight(g, v1_id, v2_id) + hfun.(Map.get(vs, v2_id))
  end

  defp do_bfs(q, %Graph{out_edges: oe} = g, target_id, %Graph{vertices: vs_tree} = tree, hfun) do
    case PriorityQueue.pop(q) do
      {{:value, {v_id, ^target_id}}, _q1} ->
        v_id_tree = Graph.Utils.vertex_id(v_id)
        construct_path(v_id_tree, tree, [target_id])
      {{:value, {v1_id, v2_id}}, q1} ->
        v2_id_tree = Graph.Utils.vertex_id(v2_id)
        if Map.has_key?(vs_tree, v2_id_tree) do
          do_bfs(q1, g, target_id, tree, hfun)
        else
          case Map.get(oe, v2_id) do
            nil ->
              do_bfs(q1, g, target_id, tree, hfun)
            v2_out ->
              tree =
                tree
                |> Graph.add_vertex(v2_id)
                |> Graph.add_edge(v2_id, v1_id)
              q2 =
                v2_out
                |> Enum.map(fn id -> {id, cost(g, v2_id, id, hfun)} end)
                |> Enum.reduce(q1, fn {id, cost}, q -> PriorityQueue.push(q, {v2_id, id}, cost) end)
              do_bfs(q2, g, target_id, tree, hfun)
          end
        end
      {:empty, _} ->
        nil
    end
  end

  defp construct_path(v_id_tree, %Graph{vertices: vs_tree, out_edges: oe_tree} = tree, path) do
    v_id_actual = Map.get(vs_tree, v_id_tree)
    path = [v_id_actual | path]
    case oe_tree |> Map.get(v_id_tree, MapSet.new) |> MapSet.to_list do
      [] ->
        path
      [next_id_tree] ->
        construct_path(next_id_tree, tree, path)
    end
  end

  defp dfs(%Graph{} = g, neighbors, target_id, path, paths) do
    paths =
      if MapSet.member?(neighbors, target_id) do
        [Enum.reverse([target_id|path])|paths]
      else
        paths
      end
    neighbors = MapSet.difference(neighbors, MapSet.new(path))
    do_dfs(g, MapSet.to_list(neighbors), target_id, path, paths)
  end

  defp do_dfs(_g, [], _target_id, _path, paths) do
    paths
  end
  defp do_dfs(%Graph{out_edges: oe} = g, [next_neighbor_id|neighbors], target_id, path, acc) do
    case Map.get(oe, next_neighbor_id) do
      nil ->
        do_dfs(g, neighbors, target_id, path, acc)
      next_neighbors ->
        case dfs(g, next_neighbors, target_id, [next_neighbor_id | path], acc) do
          nil ->
            do_dfs(g, neighbors, target_id, path, acc)
          paths ->
            do_dfs(g, neighbors, target_id, path, paths)
        end
    end
  end
end