yog_ex

defmodule Yog.Pathfinding.Dijkstra do
  @moduledoc """
  Dijkstra's algorithm for single-source shortest paths.

  Dijkstra's algorithm finds the shortest path from a source node to all other
  reachable nodes in a graph with non-negative edge weights.

  ## Algorithm Characteristics

  - **Time Complexity**: O((V + E) log V) with a binary heap
  - **Space Complexity**: O(V)
  - **Requirements**: Non-negative edge weights
  - **Optimality**: Guaranteed optimal for graphs with non-negative weights

  ## When to Use

  - When all edge weights are non-negative
  - For single-source shortest path problems
  - When you need paths to all nodes from a source
  - As a baseline comparison for other algorithms

  ## Examples

      # Find shortest path between two nodes
      graph = Yog.directed()
      |> Yog.add_node(:a, nil)
      |> Yog.add_node(:b, nil)
      |> Yog.add_node(:c, nil)
      |> Yog.add_edge!(:a, :b, 4)
      |> Yog.add_edge!(:b, :c, 1)

      compare = &Yog.Utils.compare/2
      Dijkstra.shortest_path(graph, :a, :c, 0, &(&1 + &2), compare)
      #=> {:some, {:path, [:a, :b, :c], 5}}

      # Find all distances from a source
      Dijkstra.single_source_distances(graph, :a, 0, &(&1 + &2), compare)
      #=> %{:a => 0, :b => 4, :c => 5}
  """

  alias Yog.Model
  alias Yog.Pathfinding.Path

  @typedoc "Result type for shortest path queries"
  @type path_result :: {:ok, Path.t()} | :error

  # ============================================================
  # Keyword-style API (for Pathfinding module delegation)
  # ============================================================

  @doc """
  Find shortest path using keyword options.

  ## Options

    * `:in` - The graph to search
    * `:from` - Starting node
    * `:to` - Target node
    * `:zero` - Identity value for the weight type
    * `:add` - Function to add two weights
    * `:compare` - Function to compare weights (`:lt`, `:eq`, `:gt`)

  ## Examples

      Pathfinding.shortest_path(
        in: graph,
        from: :a,
        to: :c,
        zero: 0,
        add: &(&1 + &2),
        compare: &Yog.Utils.compare/2
      )
  """
  @spec shortest_path(keyword()) :: path_result()
  def shortest_path(opts) do
    graph = Keyword.fetch!(opts, :in)
    from = Keyword.fetch!(opts, :from)
    to = Keyword.fetch!(opts, :to)
    zero = Keyword.fetch!(opts, :zero)
    add = Keyword.fetch!(opts, :add)
    compare = Keyword.fetch!(opts, :compare)

    shortest_path(graph, from, to, zero, add, compare)
  end

  @doc """
  Find shortest path using keyword options (alias for `shortest_path/1`).
  """
  @spec shortest_path_int(keyword()) :: path_result()
  def shortest_path_int(opts) do
    shortest_path(opts)
  end

  @doc """
  Find shortest path using keyword options (alias for `shortest_path/1`).
  """
  @spec shortest_path_float(keyword()) :: path_result()
  def shortest_path_float(opts) do
    shortest_path(opts)
  end

  @doc """
  Single-source distances using keyword options.

  ## Options

    * `:in` - The graph to search
    * `:from` - Starting node
    * `:zero` - Identity value for the weight type
    * `:add` - Function to add two weights
    * `:compare` - Function to compare weights

  ## Examples

      Pathfinding.single_source_distances(
        in: graph,
        from: :a,
        zero: 0,
        add: &(&1 + &2),
        compare: &Yog.Utils.compare/2
      )
  """
  @spec single_source_distances(keyword()) :: %{Yog.node_id() => any()}
  def single_source_distances(opts) do
    graph = Keyword.fetch!(opts, :in)
    from = Keyword.fetch!(opts, :from)
    zero = Keyword.fetch!(opts, :zero)
    add = Keyword.fetch!(opts, :add)
    compare = Keyword.fetch!(opts, :compare)

    single_source_distances(graph, from, zero, add, compare)
  end

  @doc """
  Implicit Dijkstra using keyword options.

  ## Options

    * `:from` - Starting state
    * `:successors_with_cost` - Function returning neighbors with costs
    * `:is_goal` - Function to check if a state is the goal
    * `:zero` - Identity value for the weight type
    * `:add` - Function to add two weights
    * `:compare` - Function to compare weights

  ## Examples

      Pathfinding.implicit_dijkstra(
        from: 1,
        successors_with_cost: fn n -> [{n+1, 1}] end,
        is_goal: fn n -> n == 10 end,
        zero: 0,
        add: &(&1 + &2),
        compare: &Yog.Utils.compare/2
      )
  """
  @spec implicit_dijkstra(keyword()) :: {:some, any()} | :none
  def implicit_dijkstra(opts) do
    from = Keyword.fetch!(opts, :from)
    successors = Keyword.fetch!(opts, :successors_with_cost)
    is_goal = Keyword.fetch!(opts, :is_goal)
    zero = Keyword.fetch!(opts, :zero)
    add = Keyword.fetch!(opts, :add)
    compare = Keyword.fetch!(opts, :compare)

    implicit_dijkstra(from, successors, is_goal, zero, add, compare)
  end

  @doc """
  Implicit Dijkstra with key function using keyword options.

  ## Options

    * `:from` - Starting state
    * `:successors_with_cost` - Function returning neighbors with costs
    * `:visited_by` - Function to extract a key for visited tracking
    * `:is_goal` - Function to check if a state is the goal
    * `:zero` - Identity value for the weight type
    * `:add` - Function to add two weights
    * `:compare` - Function to compare weights
  """
  @spec implicit_dijkstra_by(keyword()) :: {:some, any()} | :none
  def implicit_dijkstra_by(opts) do
    from = Keyword.fetch!(opts, :from)
    successors = Keyword.fetch!(opts, :successors_with_cost)
    visited_by = Keyword.fetch!(opts, :visited_by)
    is_goal = Keyword.fetch!(opts, :is_goal)
    zero = Keyword.fetch!(opts, :zero)
    add = Keyword.fetch!(opts, :add)
    compare = Keyword.fetch!(opts, :compare)

    implicit_dijkstra_by(from, successors, visited_by, is_goal, zero, add, compare)
  end

  # ============================================================
  # Direct API
  # ============================================================

  @doc """
  Find the shortest path between two nodes using custom numeric operations.

  ## Parameters

    * `graph` - The graph to search
    * `from` - Starting node
    * `to` - Target node
    * `zero` - Identity value for the weight type
    * `add` - Function to add two weights: `(weight, weight) -> weight`
    * `compare` - Function to compare weights, returns `:lt`, `:eq`, or `:gt`

  ## Returns

    * `{:ok, path}` - A `Path` struct containing the nodes and total weight
    * `:error` - No path exists between the nodes

  ## Examples

      iex> graph = Yog.directed()
      ...> |> Yog.add_node(:a, nil)
      ...> |> Yog.add_node(:b, nil)
      ...> |> Yog.add_node(:c, nil)
      ...> |> Yog.add_edge!(:a, :b, 4)
      ...> |> Yog.add_edge!(:b, :c, 1)
      iex> compare = &Yog.Utils.compare/2
      iex> {:ok, path} = Dijkstra.shortest_path(graph, :a, :c, 0, &(&1 + &2), compare)
      iex> path.nodes
      [:a, :b, :c]
      iex> path.weight
      5

      iex> graph = Yog.directed()
      ...> |> Yog.add_node(:a, nil)
      ...> |> Yog.add_node(:b, nil)
      ...> |> Yog.add_node(:c, nil)
      ...> |> Yog.add_edge!(:a, :b, 4)
      ...> |> Yog.add_edge!(:b, :c, 1)
      iex> compare = &Yog.Utils.compare/2
      iex> Dijkstra.shortest_path(graph, :a, :nonexistent, 0, &(&1 + &2), compare)
      :error
  """
  @spec shortest_path(
          Yog.t(),
          Yog.node_id(),
          Yog.node_id(),
          weight,
          (weight, weight -> weight),
          (weight, weight -> :lt | :eq | :gt)
        ) :: path_result()
        when weight: var
  def shortest_path(graph, from, to, zero, add, compare) do
    case do_dijkstra(graph, from, to, zero, add, compare, true) do
      :error -> :error
      {path, weight} -> {:ok, Path.new(path, weight, :dijkstra)}
    end
  end

  @doc """
  Find the shortest path using integer weights.

  Uses built-in integer arithmetic for efficient computation.

  ## Examples

      iex> graph = Yog.directed()
      ...> |> Yog.add_node(1, nil)
      ...> |> Yog.add_node(2, nil)
      ...> |> Yog.add_node(3, nil)
      ...> |> Yog.add_edge!(1, 2, 4)
      ...> |> Yog.add_edge!(2, 3, 1)
      iex> {:ok, path} = Dijkstra.shortest_path_int(graph, 1, 3)
      iex> path.nodes
      [1, 2, 3]
      iex> path.weight
      5
  """
  @spec shortest_path_int(Yog.t(), Yog.node_id(), Yog.node_id()) :: path_result()
  def shortest_path_int(graph, from, to) do
    shortest_path(graph, from, to, 0, &(&1 + &2), &Yog.Utils.compare/2)
  end

  @doc """
  Find the shortest path using float weights.

  Uses built-in float arithmetic for efficient computation.

  ## Examples

      iex> graph = Yog.directed()
      ...> |> Yog.add_node(1, nil)
      ...> |> Yog.add_node(2, nil)
      ...> |> Yog.add_node(3, nil)
      ...> |> Yog.add_edge!(1, 2, 4.5)
      ...> |> Yog.add_edge!(2, 3, 1.5)
      iex> {:ok, path} = Dijkstra.shortest_path_float(graph, 1, 3)
      iex> path.nodes
      [1, 2, 3]
      iex> path.weight
      6.0
  """
  @spec shortest_path_float(Yog.t(), Yog.node_id(), Yog.node_id()) :: path_result()
  def shortest_path_float(graph, from, to) do
    shortest_path(graph, from, to, 0.0, &(&1 + &2), &Yog.Utils.compare/2)
  end

  @doc """
  Calculate single-source shortest distances to all reachable nodes.

  Returns a map of node IDs to their shortest distance from the source.

  ## Parameters

    * `graph` - The graph to search
    * `from` - Source node
    * `zero` - Identity value for the weight type
    * `add` - Function to add two weights
    * `compare` - Function to compare weights

  ## Examples

      iex> graph = Yog.directed()
      ...> |> Yog.add_node(:a, nil)
      ...> |> Yog.add_node(:b, nil)
      ...> |> Yog.add_node(:c, nil)
      ...> |> Yog.add_edge!(:a, :b, 4)
      ...> |> Yog.add_edge!(:a, :c, 2)
      ...> |> Yog.add_edge!(:b, :c, 1)
      iex> compare = &Yog.Utils.compare/2
      iex> Dijkstra.single_source_distances(graph, :a, 0, &(&1 + &2), compare)
      %{a: 0, b: 4, c: 2}
  """
  @spec single_source_distances(
          Yog.t(),
          Yog.node_id(),
          weight,
          (weight, weight -> weight),
          (weight, weight -> :lt | :eq | :gt)
        ) :: %{Yog.node_id() => weight}
        when weight: var
  def single_source_distances(graph, from, zero, add, compare) do
    case do_dijkstra(graph, from, nil, zero, add, compare, false) do
      :none -> %{}
      {_path, _weight, distances} -> distances
    end
  end

  @doc """
  Run Dijkstra on an implicit (generated) graph.

  Instead of storing all edges explicitly, provide a successor function that
  generates neighbors on demand. This is useful for:
  - Infinite or very large graphs
  - Grid-based pathfinding with dynamic obstacles
  - Game state spaces

  ## Parameters

    * `from` - Starting state
    * `successors` - Function `state -> [{neighbor, cost}]`
    * `is_goal` - Function `state -> boolean` to check if goal reached
    * `zero` - Identity value for the weight type
    * `add` - Function to add two weights
    * `compare` - Function to compare weights

  ## Returns

    * `{:ok, cost}` - Minimum cost to reach goal
    * `:error` - Goal is unreachable

  ## Examples

      # Search on a linear chain: 1->2->3->4 with costs 1,2,3
      iex> successors = fn
      ...>   1 -> [{2, 1}]
      ...>   2 -> [{3, 2}]
      ...>   3 -> [{4, 3}]
      ...>   4 -> []
      ...> end
      iex> compare = &Yog.Utils.compare/2
      iex> {:ok, cost} = Dijkstra.implicit_dijkstra(
      ...>   1, successors, fn x -> x == 4 end,
      ...>   0, &(&1 + &2), compare
      ...> )
      iex> cost
      6
  """
  @spec implicit_dijkstra(
          state,
          (state -> [{state, cost}]),
          (state -> boolean),
          cost,
          (cost, cost -> cost),
          (cost, cost -> :lt | :eq | :gt)
        ) :: {:ok, cost} | :error
        when state: var, cost: var
  def implicit_dijkstra(from, successors, is_goal, zero, add, compare) do
    implicit_dijkstra_by(from, successors, fn x -> x end, is_goal, zero, add, compare)
  end

  @doc """
  Implicit Dijkstra with a key function for visited state tracking.

  Similar to `implicit_dijkstra/6`, but uses a key function to determine
  when states should be considered "visited". This allows:
  - Efficient pruning of equivalent states
  - Custom equivalence relations beyond simple equality

  ## Parameters

    * `from` - Starting state
    * `successors` - Function `state -> [{neighbor, cost}]`
    * `key_fn` - Function `state -> key` for visited tracking
    * `is_goal` - Function `state -> boolean` to check if goal reached
    * `zero` - Identity value for the weight type
    * `add` - Function to add two weights
    * `compare` - Function to compare weights

  ## Examples

      iex> successors = fn
      ...>   {pos, _dir} when pos < 3 -> [{{pos + 1, :fwd}, 1}]
      ...>   _ -> []
      ...> end
      iex> key_fn = fn {pos, _dir} -> pos end
      iex> goal_fn = fn {pos, _dir} -> pos == 3 end
      iex> compare = &Yog.Utils.compare/2
      iex> {:ok, cost} = Dijkstra.implicit_dijkstra_by(
      ...>   {0, :start}, successors, key_fn,
      ...>   goal_fn, 0, &(&1 + &2), compare
      ...> )
      iex> cost
      3
  """
  @spec implicit_dijkstra_by(
          state,
          (state -> [{state, cost}]),
          (state -> term()),
          (state -> boolean),
          cost,
          (cost, cost -> cost),
          (cost, cost -> :lt | :eq | :gt)
        ) :: {:ok, cost} | :error
        when state: var, cost: var
  def implicit_dijkstra_by(from, successors, key_fn, is_goal, zero, add, compare) do
    initial_queue = [{zero, from}]
    initial_visited = %{key_fn.(from) => zero}

    do_implicit_dijkstra(
      initial_queue,
      successors,
      key_fn,
      is_goal,
      add,
      compare,
      initial_visited
    )
  end

  # Main Dijkstra implementation
  # Returns :error | {path, weight} if return_path=true
  # Returns :error | {path, weight, distances} if return_path=false
  defp do_dijkstra(graph, from, to, zero, add, compare, return_path) do
    initial_queue = [{zero, from, [from]}]
    initial_distances = %{from => zero}

    do_dijkstra_loop(graph, initial_queue, to, add, compare, return_path, initial_distances)
  end

  defp do_dijkstra_loop(_graph, [], to, _add, _compare, _return_path, distances) do
    if to == nil do
      # Single-source case: return all distances
      {[], 0, distances}
    else
      :error
    end
  end

  defp do_dijkstra_loop(
         graph,
         [{dist, node, path} | rest],
         to,
         add,
         compare,
         return_path,
         distances
       ) do
    # Check if we've found a better path already
    current_best = Map.get(distances, node)

    if current_best != nil and compare.(dist, current_best) == :gt do
      # Skip this outdated entry
      do_dijkstra_loop(graph, rest, to, add, compare, return_path, distances)
    else
      # Check if we reached the target (only if to is specified)
      if to != nil and node == to do
        if return_path do
          {Enum.reverse(path), dist}
        else
          {Enum.reverse(path), dist, distances}
        end
      else
        # Expand neighbors
        successors = Model.successors(graph, node)

        {new_queue, new_distances} =
          Enum.reduce(successors, {rest, distances}, fn {neighbor, weight}, {q, d} ->
            new_dist = add.(dist, weight)

            case Map.fetch(d, neighbor) do
              {:ok, current} ->
                if compare.(new_dist, current) == :lt do
                  new_q = insert_sorted(q, {new_dist, neighbor, [neighbor | path]}, compare)
                  new_d = Map.put(d, neighbor, new_dist)
                  {new_q, new_d}
                else
                  {q, d}
                end

              :error ->
                new_q = insert_sorted(q, {new_dist, neighbor, [neighbor | path]}, compare)
                new_d = Map.put(d, neighbor, new_dist)
                {new_q, new_d}
            end
          end)

        if to == nil do
          # Single-source: continue until queue is empty
          do_dijkstra_loop(graph, new_queue, to, add, compare, return_path, new_distances)
        else
          # Shortest path: continue searching
          do_dijkstra_loop(graph, new_queue, to, add, compare, return_path, new_distances)
        end
      end
    end
  end

  # Implicit Dijkstra implementation
  defp do_implicit_dijkstra([], _successors, _key_fn, _is_goal, _add, _compare, _visited) do
    :error
  end

  defp do_implicit_dijkstra(
         [{dist, state} | rest],
         successors,
         key_fn,
         is_goal,
         add,
         compare,
         visited
       ) do
    key = key_fn.(state)
    current_best = Map.get(visited, key)

    if current_best != nil and compare.(dist, current_best) == :gt do
      # Skip outdated entry
      do_implicit_dijkstra(rest, successors, key_fn, is_goal, add, compare, visited)
    else
      if is_goal.(state) do
        {:ok, dist}
      else
        # Expand successors
        next_states = successors.(state)

        {new_queue, new_visited} =
          Enum.reduce(next_states, {rest, visited}, fn {next_state, cost}, {q, v} ->
            next_key = key_fn.(next_state)
            new_dist = add.(dist, cost)

            case Map.fetch(v, next_key) do
              {:ok, current} ->
                if compare.(new_dist, current) == :lt do
                  new_q = insert_sorted(q, {new_dist, next_state}, compare)
                  new_v = Map.put(v, next_key, new_dist)
                  {new_q, new_v}
                else
                  {q, v}
                end

              :error ->
                new_q = insert_sorted(q, {new_dist, next_state}, compare)
                new_v = Map.put(v, next_key, new_dist)
                {new_q, new_v}
            end
          end)

        do_implicit_dijkstra(new_queue, successors, key_fn, is_goal, add, compare, new_visited)
      end
    end
  end

  # Insert into sorted list (priority queue)
  defp insert_sorted([], item, _compare), do: [item]

  defp insert_sorted([head | tail], item, compare) do
    # Handle both 2-tuples {dist, node} and 3-tuples {dist, node, path}
    dist_item = elem(item, 0)
    dist_head = elem(head, 0)

    if compare.(dist_item, dist_head) == :lt do
      [item, head | tail]
    else
      [head | insert_sorted(tail, item, compare)]
    end
  end
end