yog

-module(yog@pathfinding@johnson).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/yog/pathfinding/johnson.gleam").
-export([johnson/5, johnson_int/1, johnson_float/1]).

-if(?OTP_RELEASE >= 27).
-define(MODULEDOC(Str), -moduledoc(Str)).
-define(DOC(Str), -doc(Str)).
-else.
-define(MODULEDOC(Str), -compile([])).
-define(DOC(Str), -compile([])).
-endif.

?MODULEDOC(
    " [Johnson's algorithm](https://en.wikipedia.org/wiki/Johnson%27s_algorithm) for\n"
    " all-pairs shortest paths in weighted graphs with negative edge weights.\n"
    "\n"
    " Johnson's algorithm efficiently computes shortest paths between all pairs of nodes\n"
    " in sparse graphs, even when edges have negative weights (but no negative cycles).\n"
    " It combines Bellman-Ford and Dijkstra's algorithms with a reweighting technique.\n"
    "\n"
    " ## Algorithm\n"
    "\n"
    " | Algorithm | Function | Complexity | Best For |\n"
    " |-----------|----------|------------|----------|\n"
    " | [Johnson's](https://en.wikipedia.org/wiki/Johnson%27s_algorithm) | `johnson/4` | O(V² log V + VE) | Sparse graphs with negative weights |\n"
    "\n"
    " ## Key Concepts\n"
    "\n"
    " - **Reweighting**: Transform negative weights to non-negative while preserving shortest paths\n"
    " - **Bellman-Ford Phase**: Compute reweighting function and detect negative cycles\n"
    " - **Dijkstra Phase**: Run Dijkstra from each vertex on reweighted graph\n"
    " - **Distance Adjustment**: Transform reweighted distances back to original weights\n"
    "\n"
    " ## How Reweighting Works\n"
    "\n"
    " The algorithm computes a potential function `h(v)` for each vertex such that:\n"
    " ```\n"
    " w'(u,v) = w(u,v) + h(u) - h(v) ≥ 0\n"
    " ```\n"
    "\n"
    " This transformation preserves shortest paths because for any path p = v₁→v₂→...→vₖ:\n"
    " ```\n"
    " w'(p) = w(p) + h(v₁) - h(vₖ)\n"
    " ```\n"
    "\n"
    " So the relative ordering of path weights remains the same!\n"
    "\n"
    " ## The Algorithm Steps\n"
    "\n"
    " 1. **Add temporary source**: Create new vertex `s` with 0-weight edges to all vertices\n"
    " 2. **Run Bellman-Ford**: From `s` to compute h(v) = distance[v] and detect negative cycles\n"
    " 3. **Reweight edges**: Set w'(u,v) = w(u,v) + h(u) - h(v) for all edges\n"
    " 4. **Run V × Dijkstra**: Compute shortest paths on reweighted graph\n"
    " 5. **Adjust distances**: Set dist(u,v) = dist'(u,v) - h(u) + h(v)\n"
    "\n"
    " ## Comparison with Other All-Pairs Algorithms\n"
    "\n"
    " | Approach | Complexity | Best For |\n"
    " |----------|------------|----------|\n"
    " | Floyd-Warshall | O(V³) | Dense graphs (E ≈ V²) |\n"
    " | Johnson's | O(V² log V + VE) | Sparse graphs (E ≪ V²) |\n"
    " | V × Dijkstra | O(V(V+E) log V) | Sparse graphs, non-negative weights only |\n"
    " | V × Bellman-Ford | O(V²E) | Rarely optimal |\n"
    "\n"
    " **Rule of thumb**:\n"
    " - Use Johnson's for sparse graphs with negative weights\n"
    " - Use Floyd-Warshall for dense graphs or when simplicity matters\n"
    " - Use V × Dijkstra for sparse graphs with only non-negative weights\n"
    "\n"
    " ## Complexity Analysis\n"
    "\n"
    " - **Bellman-Ford phase**: O(VE) - run once\n"
    " - **Dijkstra phase**: O(V × (V+E) log V) = O(V² log V + VE) - run V times\n"
    " - **Total**: O(V² log V + VE)\n"
    "\n"
    " For sparse graphs where E = O(V), this is O(V² log V), much better than\n"
    " Floyd-Warshall's O(V³)!\n"
    "\n"
    " ## Negative Cycles\n"
    "\n"
    " The algorithm detects negative cycles during the Bellman-Ford phase.\n"
    " If a negative cycle exists, the function returns `Error(Nil)`.\n"
    "\n"
    " ## Use Cases\n"
    "\n"
    " - **Sparse road networks**: Computing all-pairs distances with tolls/credits\n"
    " - **Currency arbitrage**: Finding profitable exchange cycles\n"
    " - **Network routing**: Precomputing routing tables with various costs\n"
    " - **Game AI**: Pathfinding in large sparse maps with varied terrain costs\n"
    "\n"
    " ## History\n"
    "\n"
    " Published by Donald B. Johnson in 1977. The algorithm is a brilliant\n"
    " combination of Bellman-Ford's ability to handle negative weights and\n"
    " Dijkstra's efficiency with non-negative weights.\n"
    "\n"
    " ## References\n"
    "\n"
    " - [Wikipedia: Johnson's Algorithm](https://en.wikipedia.org/wiki/Johnson%27s_algorithm)\n"
    " - [Johnson's Original Paper (1977)](https://doi.org/10.1145/321992.321993)\n"
    " - [CP-Algorithms: Johnson's Algorithm](https://cp-algorithms.com/graph/all-pair-shortest-path-johnson.html)\n"
    " - [MIT 6.006: Advanced Shortest Paths](https://ocw.mit.edu/courses/6-006-introduction-to-algorithms-spring-2020/)\n"
).

-file("src/yog/pathfinding/johnson.gleam", 183).
?DOC(" Finds a node ID that doesn't exist in the graph (for temporary source)\n").
-spec find_new_node_id(list(integer())) -> integer().
find_new_node_id(Nodes) ->
    case Nodes of
        [] ->
            0;

        _ ->
            Max_node = gleam@list:fold(
                Nodes,
                0,
                fun(Max, Node) -> case Node > Max of
                        true ->
                            Node;

                        false ->
                            Max
                    end end
            ),
            Max_node + 1
    end.

-file("src/yog/pathfinding/johnson.gleam", 246).
?DOC(" Modified Bellman-Ford relaxation that includes edges from temporary source\n").
-spec bellman_ford_with_temp_source(
    yog@model:graph(any(), YRF),
    integer(),
    list(integer()),
    gleam@dict:dict(integer(), YRF),
    gleam@dict:dict(integer(), integer()),
    integer(),
    YRF,
    fun((YRF, YRF) -> YRF),
    fun((YRF, YRF) -> gleam@order:order())
) -> {gleam@dict:dict(integer(), YRF), gleam@dict:dict(integer(), integer())}.
bellman_ford_with_temp_source(
    Graph,
    Temp_source,
    All_nodes,
    Distances,
    Predecessors,
    Remaining,
    Zero,
    Add,
    Compare
) ->
    case Remaining =< 0 of
        true ->
            {Distances, Predecessors};

        false ->
            {New_distances, New_predecessors} = gleam@list:fold(
                All_nodes,
                {Distances, Predecessors},
                fun(Acc, U) ->
                    {Dists, Preds} = Acc,
                    case gleam_stdlib:map_get(Dists, U) of
                        {error, nil} ->
                            Acc;

                        {ok, U_dist} ->
                            Neighbors = case U =:= Temp_source of
                                true ->
                                    _pipe = gleam@list:filter(
                                        All_nodes,
                                        fun(Node) -> Node /= Temp_source end
                                    ),
                                    gleam@list:map(
                                        _pipe,
                                        fun(Node@1) -> {Node@1, Zero} end
                                    );

                                false ->
                                    yog@model:successors(Graph, U)
                            end,
                            gleam@list:fold(
                                Neighbors,
                                {Dists, Preds},
                                fun(Inner_acc, Edge) ->
                                    {V, Weight} = Edge,
                                    {Curr_dists, Curr_preds} = Inner_acc,
                                    New_dist = Add(U_dist, Weight),
                                    case gleam_stdlib:map_get(Curr_dists, V) of
                                        {error, nil} ->
                                            {gleam@dict:insert(
                                                    Curr_dists,
                                                    V,
                                                    New_dist
                                                ),
                                                gleam@dict:insert(
                                                    Curr_preds,
                                                    V,
                                                    U
                                                )};

                                        {ok, V_dist} ->
                                            case Compare(New_dist, V_dist) of
                                                lt ->
                                                    {gleam@dict:insert(
                                                            Curr_dists,
                                                            V,
                                                            New_dist
                                                        ),
                                                        gleam@dict:insert(
                                                            Curr_preds,
                                                            V,
                                                            U
                                                        )};

                                                _ ->
                                                    Inner_acc
                                            end
                                    end
                                end
                            )
                    end
                end
            ),
            bellman_ford_with_temp_source(
                Graph,
                Temp_source,
                All_nodes,
                New_distances,
                New_predecessors,
                Remaining - 1,
                Zero,
                Add,
                Compare
            )
    end.

-file("src/yog/pathfinding/johnson.gleam", 315).
?DOC(" Checks for negative cycles including edges from temporary source\n").
-spec has_negative_cycle_with_temp_source(
    yog@model:graph(any(), YRS),
    integer(),
    list(integer()),
    gleam@dict:dict(integer(), YRS),
    YRS,
    fun((YRS, YRS) -> YRS),
    fun((YRS, YRS) -> gleam@order:order())
) -> boolean().
has_negative_cycle_with_temp_source(
    Graph,
    Temp_source,
    Nodes,
    Distances,
    Zero,
    Add,
    Compare
) ->
    All_nodes = [Temp_source | Nodes],
    gleam@list:any(
        All_nodes,
        fun(U) -> case gleam_stdlib:map_get(Distances, U) of
                {error, nil} ->
                    false;

                {ok, U_dist} ->
                    Neighbors = case U =:= Temp_source of
                        true ->
                            _pipe = gleam@list:filter(
                                All_nodes,
                                fun(Node) -> Node /= Temp_source end
                            ),
                            gleam@list:map(
                                _pipe,
                                fun(Node@1) -> {Node@1, Zero} end
                            );

                        false ->
                            yog@model:successors(Graph, U)
                    end,
                    gleam@list:any(
                        Neighbors,
                        fun(Edge) ->
                            {V, Weight} = Edge,
                            New_dist = Add(U_dist, Weight),
                            case gleam_stdlib:map_get(Distances, V) of
                                {error, nil} ->
                                    false;

                                {ok, V_dist} ->
                                    case Compare(New_dist, V_dist) of
                                        lt ->
                                            true;

                                        _ ->
                                            false
                                    end
                            end
                        end
                    )
            end end
    ).

-file("src/yog/pathfinding/johnson.gleam", 200).
?DOC(" Runs Bellman-Ford to compute reweighting function h(v) for each vertex\n").
-spec compute_reweighting_function(
    yog@model:graph(any(), YQW),
    integer(),
    list(integer()),
    YQW,
    fun((YQW, YQW) -> YQW),
    fun((YQW, YQW) -> gleam@order:order())
) -> {ok, gleam@dict:dict(integer(), YQW)} | {error, nil}.
compute_reweighting_function(Graph, Temp_source, Nodes, Zero, Add, Compare) ->
    Initial_distances = maps:from_list([{Temp_source, Zero}]),
    Initial_predecessors = maps:new(),
    Node_count = erlang:length(Nodes) + 1,
    {Distances, _} = bellman_ford_with_temp_source(
        Graph,
        Temp_source,
        [Temp_source | Nodes],
        Initial_distances,
        Initial_predecessors,
        Node_count - 1,
        Zero,
        Add,
        Compare
    ),
    case has_negative_cycle_with_temp_source(
        Graph,
        Temp_source,
        Nodes,
        Distances,
        Zero,
        Add,
        Compare
    ) of
        true ->
            {error, nil};

        false ->
            {ok, Distances}
    end.

-file("src/yog/pathfinding/johnson.gleam", 393).
?DOC(" Helper to unwrap Result with a provided default value\n").
-spec unwrap_or({ok, YSG} | {error, any()}, YSG) -> YSG.
unwrap_or(Result, Default) ->
    case Result of
        {ok, Value} ->
            Value;

        {error, _} ->
            Default
    end.

-file("src/yog/pathfinding/johnson.gleam", 356).
?DOC(" Creates a new graph with reweighted edges: w'(u,v) = w(u,v) + h(u) - h(v)\n").
-spec reweight_graph(
    yog@model:graph(YRY, YRZ),
    gleam@dict:dict(integer(), YRZ),
    YRZ,
    fun((YRZ, YRZ) -> YRZ),
    fun((YRZ, YRZ) -> YRZ)
) -> yog@model:graph(YRY, YRZ).
reweight_graph(Graph, H, Zero, Add, Subtract) ->
    New_out_edges = gleam@dict:fold(
        erlang:element(4, Graph),
        maps:new(),
        fun(Acc, U, Neighbors) ->
            H_u = begin
                _pipe = gleam_stdlib:map_get(H, U),
                unwrap_or(_pipe, Zero)
            end,
            New_neighbors = gleam@dict:fold(
                Neighbors,
                maps:new(),
                fun(Inner_acc, V, Weight) ->
                    H_v = begin
                        _pipe@1 = gleam_stdlib:map_get(H, V),
                        unwrap_or(_pipe@1, Zero)
                    end,
                    New_weight = Subtract(Add(Weight, H_u), H_v),
                    gleam@dict:insert(Inner_acc, V, New_weight)
                end
            ),
            gleam@dict:insert(Acc, U, New_neighbors)
        end
    ),
    New_in_edges = gleam@dict:fold(
        erlang:element(5, Graph),
        maps:new(),
        fun(Acc@1, V@1, Sources) ->
            H_v@1 = begin
                _pipe@2 = gleam_stdlib:map_get(H, V@1),
                unwrap_or(_pipe@2, Zero)
            end,
            New_sources = gleam@dict:fold(
                Sources,
                maps:new(),
                fun(Inner_acc@1, U@1, Weight@1) ->
                    H_u@1 = begin
                        _pipe@3 = gleam_stdlib:map_get(H, U@1),
                        unwrap_or(_pipe@3, Zero)
                    end,
                    New_weight@1 = Subtract(Add(Weight@1, H_u@1), H_v@1),
                    gleam@dict:insert(Inner_acc@1, U@1, New_weight@1)
                end
            ),
            gleam@dict:insert(Acc@1, V@1, New_sources)
        end
    ),
    {graph,
        erlang:element(2, Graph),
        erlang:element(3, Graph),
        New_out_edges,
        New_in_edges}.

-file("src/yog/pathfinding/johnson.gleam", 131).
?DOC(
    " Computes shortest paths between all pairs of nodes using Johnson's algorithm.\n"
    "\n"
    " **Time Complexity:** O(V² log V + VE)\n"
    "\n"
    " Returns an error if a negative cycle is detected.\n"
    "\n"
    " ## Parameters\n"
    "\n"
    " - `zero`: The identity element for addition (e.g., 0 for integers)\n"
    " - `add`: Function to add two weights\n"
    " - `subtract`: Function to subtract two weights (for reweighting)\n"
    " - `compare`: Function to compare two weights\n"
    "\n"
    " ## Example\n"
    "\n"
    " ```gleam\n"
    " let result = johnson.johnson(\n"
    "   in: graph,\n"
    "   with_zero: 0,\n"
    "   with_add: int.add,\n"
    "   with_subtract: int.subtract,\n"
    "   with_compare: int.compare\n"
    " )\n"
    " // => Ok(Dict([#(#(1, 2), 10), #(#(1, 3), 25), ...]))\n"
    " ```\n"
    "\n"
    " ## When to Use\n"
    "\n"
    " Use Johnson's algorithm for **sparse graphs** when you need all-pairs shortest\n"
    " paths and the graph may contain negative edge weights (but no negative cycles).\n"
    " For dense graphs, prefer Floyd-Warshall. For graphs with only non-negative\n"
    " weights, consider running Dijkstra from each vertex directly.\n"
).
-spec johnson(
    yog@model:graph(any(), YQN),
    YQN,
    fun((YQN, YQN) -> YQN),
    fun((YQN, YQN) -> YQN),
    fun((YQN, YQN) -> gleam@order:order())
) -> {ok, gleam@dict:dict({integer(), integer()}, YQN)} | {error, nil}.
johnson(Graph, Zero, Add, Subtract, Compare) ->
    Nodes = yog@model:all_nodes(Graph),
    Temp_source = find_new_node_id(Nodes),
    case compute_reweighting_function(
        Graph,
        Temp_source,
        Nodes,
        Zero,
        Add,
        Compare
    ) of
        {error, nil} ->
            {error, nil};

        {ok, H} ->
            All_pairs_distances = gleam@list:fold(
                Nodes,
                maps:new(),
                fun(Distances, U) ->
                    Distances_from_u = yog@pathfinding@dijkstra:single_source_distances(
                        reweight_graph(Graph, H, Zero, Add, Subtract),
                        U,
                        Zero,
                        Add,
                        Compare
                    ),
                    gleam@list:fold(
                        Nodes,
                        Distances,
                        fun(Acc, V) ->
                            case gleam_stdlib:map_get(Distances_from_u, V) of
                                {ok, Dist_reweighted} ->
                                    H_u = begin
                                        _pipe = gleam_stdlib:map_get(H, U),
                                        unwrap_or(_pipe, Zero)
                                    end,
                                    H_v = begin
                                        _pipe@1 = gleam_stdlib:map_get(H, V),
                                        unwrap_or(_pipe@1, Zero)
                                    end,
                                    Actual_dist = Add(
                                        Subtract(Dist_reweighted, H_u),
                                        H_v
                                    ),
                                    gleam@dict:insert(Acc, {U, V}, Actual_dist);

                                {error, nil} ->
                                    Acc
                            end
                        end
                    )
                end
            ),
            {ok, All_pairs_distances}
    end.

-file("src/yog/pathfinding/johnson.gleam", 426).
?DOC(
    " Computes all-pairs shortest paths with **integer weights** using Johnson's algorithm.\n"
    "\n"
    " This is a convenience wrapper around `johnson` that uses:\n"
    " - `0` as the zero element\n"
    " - `int.add` for addition\n"
    " - `int.subtract` for subtraction\n"
    " - `int.compare` for comparison\n"
    "\n"
    " ## Example\n"
    "\n"
    " ```gleam\n"
    " let result = johnson.johnson_int(graph)\n"
    " // => Ok(Dict([#(#(1, 2), 10), #(#(1, 3), 25), ...]))\n"
    " ```\n"
    "\n"
    " ## When to Use\n"
    "\n"
    " Use this for **sparse graphs** where you need all-pairs distances with `Int`\n"
    " weights that may be negative. For dense graphs, prefer `floyd_warshall_int`.\n"
    " For graphs with only non-negative weights, consider running `dijkstra` V times.\n"
    "\n"
    " Returns `Error(Nil)` if a negative cycle is detected.\n"
).
-spec johnson_int(yog@model:graph(any(), integer())) -> {ok,
        gleam@dict:dict({integer(), integer()}, integer())} |
    {error, nil}.
johnson_int(Graph) ->
    johnson(
        Graph,
        0,
        fun gleam@int:add/2,
        fun gleam@int:subtract/2,
        fun gleam@int:compare/2
    ).

-file("src/yog/pathfinding/johnson.gleam", 451).
?DOC(
    " Computes all-pairs shortest paths with **float weights** using Johnson's algorithm.\n"
    "\n"
    " This is a convenience wrapper around `johnson` that uses:\n"
    " - `0.0` as the zero element\n"
    " - `float.add` for addition\n"
    " - `float.subtract` for subtraction\n"
    " - `float.compare` for comparison\n"
    "\n"
    " ## Warning\n"
    "\n"
    " Float arithmetic has precision limitations. Negative cycles might not be\n"
    " detected reliably due to floating-point errors. Prefer `Int` weights for\n"
    " critical calculations.\n"
).
-spec johnson_float(yog@model:graph(any(), float())) -> {ok,
        gleam@dict:dict({integer(), integer()}, float())} |
    {error, nil}.
johnson_float(Graph) ->
    johnson(
        Graph,
        +0.0,
        fun gleam@float:add/2,
        fun gleam@float:subtract/2,
        fun gleam@float:compare/2
    ).