defmodule Scurry.Astar do
  @moduledoc """

  Implementation of [A-star search
  algorithm](https://en.wikipedia.org/wiki/A*_search_algorithm) to find the
  shortest path in graphs or 2D polygon maps.

  Basic usage is;

  ```
  # Define function for the heuristic cost between two nodes
  def heur_fun(node_a, node_b) do
    ... returns the cost of node a to node b
  end

  # Define a graph
  graph = %{
    node_1 => [
      {node_2, cost_1_to_2}, {node_3, cost_1_to_3},
    ],
    node_2 => [
      {node_3, cost_2_to_3}, {node_4, cost_2_to_4},
    ],
    ...
  }

  # Do A-star search
  state = Astar.search(graph, node_1, node_z, &heur_fun/2)

  # Extract path from the state
  path = Astar.path(state)
  [node1, ..., node_z]
  ```

  See [Quickstart](quickstart.html) for a concrete end-to-end example of
  defining a world and holes/obstacles, then using `PolygonMap` and `Astar`
  modules to do path finding within this.
  """

  use Scurry.Astar.Types

  @typedoc "The internal state of the A-star algorithm. Use `path/1` to extract the result."
  @opaque state :: map()

  @doc """
  Find shortest path in `graph` from `start` to `stop`.

  ## Params

  * `graph` a map of node to a list of nodes and cost tuples.
  * `start` the node from which to start
  * `stop` the node at which to end
  * `heur_fun` the heuristic function used to estimate cost from a node in
    `graph` to `stop`. It takes two nodes and returns a cost that should be
    comparable with itself for ordering. `node, node :: term`.

  ## Returns

  The algorithms internal state which can be passed to `path/1` to
  obtain the actual path.
  """
  @spec search(graph(), gnode(), gnode(), cost_fun()) :: state()
  def search(graph, start, stop, heur_fun) do
    queue = [start]

    %{
      start: start,
      stop: stop,
      graph: graph,

      # (node, node) :: cost function
      heur_fun: heur_fun,
      queue: queue,
      shortest_path_tree: %{},
      frontier: %{},

      # node => cost - distance from start to node
      g_cost: %{},
      # vertice => cost, distance from start to vervice + heuristic distance to stop
      f_cost: %{}
    }
    |> pathfind_helper
  end

  defp sort_queue(queue, f_cost) do
    Enum.sort_by(queue, fn e -> Map.get(f_cost, e) end, :asc)
  end

  defp add_to_queue(queue, node) do
    Enum.sort(queue ++ [node])
    |> Enum.dedup()
  end

  # Queue exhausted without reaching stop - stop is unreachable from start.
  defp pathfind_helper(%{queue: []} = state) do
    state
  end

  defp pathfind_helper(%{queue: [current | queue]} = state) do
    # spt is Shortest Path Tree
    spt = Map.put(state.shortest_path_tree, current, Map.get(state.frontier, current))

    cond do
      current == state.stop ->
        %{state | shortest_path_tree: spt}

      true ->
        edges = Map.get(state.graph, current, [])

        reduce_seed = {state.frontier, queue, state.g_cost, state.f_cost}

        {f, q, g_cost, f_cost} =
          Enum.reduce(edges, reduce_seed, fn {node, edge_cost}, acc ->
            {frontier, queue, g_cost, f_cost} = acc
            # H cost, the heuristic cost to get to this node
            heur_cost = state.heur_fun.(node, state.stop)
            # G cost, the current shortest cost from the start node
            shortest_distance_from_start = Map.get(g_cost, current, 0) + edge_cost
            # F cost = G cost + H cost, the cost to get to this node
            total_distance = shortest_distance_from_start + heur_cost

            cond do
              node == state.start ->
                # Never go back to the start node (cycle)
                acc

              not Map.has_key?(frontier, node) ->
                # We have not visited this node before, so account for it's cost.
                {
                  Map.put(frontier, node, current),
                  add_to_queue(queue, node),
                  Map.put(g_cost, node, shortest_distance_from_start),
                  Map.put(f_cost, node, total_distance)
                }

              shortest_distance_from_start < Map.get(g_cost, node, 0) and
                Map.get(spt, node) == nil ->
                # We've already seen this node and it's still open (not yet
                # finalized in spt), but this path to it is cheaper than what
                # we had — update its cost. It's already in the queue from when
                # it was first discovered, so no need to add_to_queue again.
                {
                  Map.put(frontier, node, current),
                  queue,
                  Map.put(g_cost, node, shortest_distance_from_start),
                  Map.put(f_cost, node, total_distance)
                }

              true ->
                {frontier, queue, g_cost, f_cost}
            end
          end)

        new_state = %{
          state
          | queue: sort_queue(q, f_cost),
            frontier: f,
            f_cost: f_cost,
            g_cost: g_cost,
            shortest_path_tree: spt
        }

        pathfind_helper(new_state)
    end
  end

  @doc """
  Get the path from the state returned by `search/4`.

  ## Params

  * `state` the a-star state returned by `search/4`

  ## Returns

  The path as a list of nodes. The type of the node is the same as in the
  `t:graph/0` used in the call to `search/4`.

  """
  @spec path(state()) :: list(gnode())
  def path(state) do
    next = state.shortest_path_tree[state.stop]

    path(state, state.start, next, [state.stop])
    |> Enum.reverse()
  end

  defp path(_state, _start, nil, acc) do
    acc
  end

  defp path(_state, start, start, acc) do
    acc ++ [start]
  end

  defp path(state, start, node, acc) do
    next = state.shortest_path_tree[node]
    path(state, start, next, acc ++ [node])
  end
end