defmodule Yog.Pathfinding.Johnson do
  @moduledoc """
  [Johnson's algorithm](https://en.wikipedia.org/wiki/Johnson%27s_algorithm) for
  all-pairs shortest paths in weighted graphs with negative edge weights.

  Johnson's algorithm efficiently computes shortest paths between all pairs of nodes
  in sparse graphs, even when edges have negative weights (but no negative cycles).
  It combines Bellman-Ford and Dijkstra's algorithms with a reweighting technique.

  ## Algorithm

  | Algorithm | Function | Complexity | Best For |
  |-----------|----------|------------|----------|
  | [Johnson's](https://en.wikipedia.org/wiki/Johnson%27s_algorithm) | `johnson/5` | O(V² log V + VE) | Sparse graphs with negative weights |

  ## Key Concepts

  - **Reweighting**: Transform negative weights to non-negative while preserving shortest paths
  - **Bellman-Ford Phase**: Compute reweighting function and detect negative cycles
  - **Dijkstra Phase**: Run Dijkstra from each vertex on reweighted graph
  - **Distance Adjustment**: Transform reweighted distances back to original weights

  ## How Reweighting Works

  The algorithm computes a potential function `h(v)` for each vertex such that:
  ```
  w'(u,v) = w(u,v) + h(u) - h(v) ≥ 0
  ```

  This transformation preserves shortest paths because for any path p = v₁→v₂→...→vₖ:
  ```
  w'(p) = w(p) + h(v₁) - h(vₖ)
  ```

  So the relative ordering of path weights remains the same!

  ## The Algorithm Steps

  1. **Add temporary source**: Create new vertex `s` with 0-weight edges to all vertices
  2. **Run Bellman-Ford**: From `s` to compute h(v) = distance[v] and detect negative cycles
  3. **Reweight edges**: Set w'(u,v) = w(u,v) + h(u) - h(v) for all edges
  4. **Run V × Dijkstra**: Compute shortest paths on reweighted graph
  5. **Adjust distances**: Set dist(u,v) = dist'(u,v) - h(u) + h(v)

  ## Comparison with Other All-Pairs Algorithms

  | Approach | Complexity | Best For |
  |----------|------------|----------|
  | Floyd-Warshall | O(V³) | Dense graphs (E ≈ V²) |
  | Johnson's | O(V² log V + VE) | Sparse graphs (E ≪ V²) |
  | V × Dijkstra | O(V(V+E) log V) | Sparse graphs, non-negative weights only |

  **Rule of thumb**:
  - Use Johnson's for sparse graphs with negative weights
  - Use Floyd-Warshall for dense graphs or when simplicity matters
  - Use V × Dijkstra for sparse graphs with only non-negative weights

  ## Negative Cycles

  The algorithm detects negative cycles during the Bellman-Ford phase.
  If a negative cycle exists, the function returns `{:error, :negative_cycle}`.

  ## Use Cases

  - **Sparse road networks**: Computing all-pairs distances with tolls/credits
  - **Currency arbitrage**: Finding profitable exchange cycles
  - **Network routing**: Precomputing routing tables with various costs
  - **Game AI**: Pathfinding in large sparse maps with varied terrain costs

  ## History

  Published by Donald B. Johnson in 1977.

  ## References

  - [Wikipedia: Johnson's Algorithm](https://en.wikipedia.org/wiki/Johnson%27s_algorithm)
  """

  alias Yog.Model
  alias Yog.Pathfinding.Dijkstra
  alias Yog.Transform

  @typedoc """
  Distance matrix: map from `{from, to}` tuple to distance.
  """
  @type distance_matrix :: %{{Yog.node_id(), Yog.node_id()} => any()}

  @doc """
  Computes all-pairs shortest paths using Johnson's algorithm.

  **Time Complexity:** O(V² log V + VE)

  Returns `{:ok, distance_matrix}` on success, or `{:error, :negative_cycle}`
  if a negative cycle is detected.

  ## Parameters

  - `graph` - The graph to analyze
  - `zero` - Identity element for addition
  - `add` - Function to add two weights
  - `subtract` - Function to subtract two weights
  - `compare` - Function to compare two weights

  ## Examples

      # Graph with negative weights (but no negative cycles)
      iex> graph = Yog.directed()
      ...> |> Yog.add_node(1, nil)
      ...> |> Yog.add_node(2, nil)
      ...> |> Yog.add_node(3, nil)
      ...> |> Yog.add_edge_ensure(from: 1, to: 2, with: 4)
      ...> |> Yog.add_edge_ensure(from: 2, to: 3, with: -3)
      ...> |> Yog.add_edge_ensure(from: 1, to: 3, with: 10)
      iex> compare = &Yog.Utils.compare/2
      iex> {:ok, distances} = Yog.Pathfinding.Johnson.johnson(graph, 0, &(&1 + &2), &(&1 - &2), compare)
      iex> # Shortest path from 1 to 3 should be 1->2->3 = 1, not direct 10
      ...> distances[{1, 3}]
      1

      # Negative cycle detection
      iex> bad_graph = Yog.directed()
      ...> |> Yog.add_node(1, nil)
      ...> |> Yog.add_node(2, nil)
      ...> |> Yog.add_edge_ensure(from: 1, to: 2, with: 1)
      ...> |> Yog.add_edge_ensure(from: 2, to: 1, with: -3)
      iex> compare = &Yog.Utils.compare/2
      iex> Yog.Pathfinding.Johnson.johnson(bad_graph, 0, &(&1 + &2), &(&1 - &2), compare)
      {:error, :negative_cycle}
  """
  @spec johnson(
          Yog.graph(),
          any(),
          (any(), any() -> any()),
          (any(), any() -> any()),
          (any(), any() -> :lt | :eq | :gt)
        ) :: {:ok, distance_matrix()} | {:error, :negative_cycle}
  def johnson(
        graph,
        zero \\ 0,
        add \\ &Kernel.+/2,
        subtract \\ &Kernel.-/2,
        compare \\ &Yog.Utils.compare/2
      ) do
    nodes = Model.all_nodes(graph)

    # Step 1 & 2: Run Bellman-Ford from temporary source to get potentials h(v)
    case compute_potentials(graph, nodes, zero, add, compare) do
      {:ok, potentials} ->
        # Step 3, 4, 5: Run Dijkstra from each node with reweighted edges
        distances = run_dijkstra_from_all(graph, nodes, potentials, zero, add, subtract, compare)
        {:ok, distances}

      {:error, :negative_cycle} ->
        {:error, :negative_cycle}
    end
  end

  # Step 1 & 2: Compute potentials h(v) using Bellman-Ford from temporary source
  defp compute_potentials(graph, nodes, zero, add, compare) do
    # Create temporary source with 0-weight edges to all nodes
    # Using make_ref() ensures no collision with user node IDs
    temp_source = make_ref()

    # Initialize distances: temp_source = 0, others = nil (infinity)
    initial_distances = %{temp_source => zero}

    # Get all edges including from temporary source
    edges = get_all_edges_with_temp_source(graph, nodes, temp_source, zero)

    # Relax edges |V| times (where |V| includes temp_source, so length(nodes) + 1)
    node_count = length(nodes) + 1

    distances =
      Enum.reduce(1..node_count, initial_distances, fn iteration, dist ->
        new_dist = relax_all_edges(edges, dist, add, compare)

        # Check for negative cycle on last iteration
        if iteration == node_count do
          if distances_changed?(dist, new_dist, compare) do
            :negative_cycle
          else
            new_dist
          end
        else
          new_dist
        end
      end)

    case distances do
      :negative_cycle -> {:error, :negative_cycle}
      dists -> {:ok, Map.delete(dists, temp_source)}
    end
  end

  # Get all edges from graph plus 0-weight edges from temp source to all nodes
  defp get_all_edges_with_temp_source(graph, nodes, temp_source, zero) do
    # Edges from temporary source to all nodes (weight = 0)
    temp_edges = Enum.map(nodes, fn node -> {temp_source, node, zero} end)

    # Regular edges from graph
    regular_edges =
      Enum.flat_map(nodes, fn u ->
        successors = Model.successors(graph, u)
        Enum.map(successors, fn {v, weight} -> {u, v, weight} end)
      end)

    temp_edges ++ regular_edges
  end

  # Relax all edges once
  defp relax_all_edges(edges, distances, add, compare) do
    Enum.reduce(edges, distances, fn {u, v, weight}, dist ->
      case Map.fetch(dist, u) do
        {:ok, dist_u} ->
          new_dist_v = add.(dist_u, weight)

          case Map.fetch(dist, v) do
            {:ok, current_dist_v} ->
              if compare.(new_dist_v, current_dist_v) == :lt do
                Map.put(dist, v, new_dist_v)
              else
                dist
              end

            :error ->
              Map.put(dist, v, new_dist_v)
          end

        :error ->
          dist
      end
    end)
  end

  # Check if distances changed (for negative cycle detection)
  defp distances_changed?(old_dist, new_dist, compare) do
    Enum.any?(new_dist, fn {node, new_val} ->
      case Map.fetch(old_dist, node) do
        {:ok, old_val} -> compare.(new_val, old_val) == :lt
        :error -> true
      end
    end)
  end

  # Run Dijkstra from each node with reweighted edges
  defp run_dijkstra_from_all(graph, nodes, potentials, zero, add, subtract, compare) do
    # Reweight the graph ONCE using map_edges_indexed
    reweighted_graph =
      Transform.map_edges_indexed(graph, fn u, v, w ->
        h_u = Map.get(potentials, u, zero)
        h_v = Map.get(potentials, v, zero)
        # w'(u,v) = w(u,v) + h(u) - h(v)
        add.(w, h_u) |> subtract.(h_v)
      end)

    parallel_opts = [
      max_concurrency: System.schedulers_online(),
      timeout: :infinity
    ]

    nodes
    |> Task.async_stream(
      fn source ->
        # Run Dijkstra on reweighted graph using the standard implementation
        reweighted_distances =
          Dijkstra.single_source_distances(reweighted_graph, source, zero, add, compare)

        # Adjust distances back: dist(u,v) = dist'(u,v) - h(u) + h(v)
        h_source = Map.get(potentials, source, zero)

        Enum.reduce(reweighted_distances, %{}, fn {dest, dist_prime}, inner_acc ->
          h_dest = Map.get(potentials, dest, zero)
          # dist = dist' - h(u) + h(v)
          adjusted_dist = add.(subtract.(dist_prime, h_source), h_dest)
          Map.put(inner_acc, {source, dest}, adjusted_dist)
        end)
      end,
      parallel_opts
    )
    |> Enum.reduce(%{}, fn {:ok, source_results}, acc ->
      Map.merge(acc, source_results)
    end)
  end
end