-module(viva_math@autodiff_reverse).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_math/autodiff_reverse.gleam").
-export([empty_tape/0, input/2, value/2, add/3, sub/3, mul/3, 'div'/3, neg/2, scale/3, exp/2, ln/2, sin/2, cos/2, tanh/2, sigmoid/2, pow/3, backward/2, grad_of/2, gradients/2]).
-export_type([op/0, node_/0, tape/0]).

-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(
    " Reverse-mode automatic differentiation via a computation tape.\n"
    "\n"
    " Forward-mode AD (`viva_math/autodiff`) is optimal when the input\n"
    " dimension is small relative to the output (e.g. one-variable\n"
    " gradients). Reverse-mode AD (\"backprop\") is optimal for the opposite\n"
    " case: scalar output, high-dimensional input — exactly the regime of\n"
    " neural networks and gradient-based MAP inference.\n"
    "\n"
    " ## How it works\n"
    "\n"
    " 1. The forward pass builds a directed acyclic graph (the \"tape\") that\n"
    "    records every elementary operation and its operands.\n"
    " 2. The backward pass walks the tape in reverse, applying the chain\n"
    "    rule to accumulate `∂output/∂node` for every node.\n"
    "\n"
    " All inputs that share a tape are differentiated in a single backward\n"
    " pass, regardless of how many there are — that's why reverse-mode AD is\n"
    " O(1) in input dimension for gradient computation.\n"
    "\n"
    " ## Limitations of this implementation\n"
    "\n"
    " - Single-output, scalar gradients only. For Jacobians of vector-valued\n"
    "   functions, use `viva_math/autodiff.jacobian` (forward) or call this\n"
    "   multiple times.\n"
    " - The tape is a `Dict(Int, ...)` for clarity; production tools use\n"
    "   linear arrays for cache locality. Adequate for graphs ≤ 10⁴ nodes.\n"
    "\n"
    " ## Example\n"
    "\n"
    " ```gleam\n"
    " import viva_math/autodiff_reverse as ad\n"
    "\n"
    " // f(x, y, z) = x² + 2·y·z\n"
    " let tape = ad.empty_tape()\n"
    " let #(x, tape) = ad.input(tape, 1.0)\n"
    " let #(y, tape) = ad.input(tape, 2.0)\n"
    " let #(z, tape) = ad.input(tape, 3.0)\n"
    " let #(x_sq, tape) = ad.mul(tape, x, x)\n"
    " let #(yz, tape) = ad.mul(tape, y, z)\n"
    " let #(yz2, tape) = ad.scale(tape, yz, 2.0)\n"
    " let #(out, tape) = ad.add(tape, x_sq, yz2)\n"
    " let grads = ad.backward(tape, out)\n"
    " // grads[x] = 2x = 2.0\n"
    " // grads[y] = 2z = 6.0\n"
    " // grads[z] = 2y = 4.0\n"
    " ```\n"
).

-type op() :: input |
    {add, integer(), integer()} |
    {sub, integer(), integer()} |
    {mul, integer(), integer()} |
    {'div', integer(), integer()} |
    {neg, integer()} |
    {scale, integer(), float()} |
    {exp, integer()} |
    {ln, integer()} |
    {sin, integer()} |
    {cos, integer()} |
    {tanh, integer()} |
    {sigmoid, integer()} |
    {pow, integer(), float()}.

-type node_() :: {node, float(), op()}.

-type tape() :: {tape, gleam@dict:dict(integer(), node_()), integer()}.

-file("src/viva_math/autodiff_reverse.gleam", 88).
-spec empty_tape() -> tape().
empty_tape() ->
    {tape, maps:new(), 0}.

-file("src/viva_math/autodiff_reverse.gleam", 92).
-spec push(tape(), float(), op()) -> {integer(), tape()}.
push(Tape, Value, Op) ->
    Id = erlang:element(3, Tape),
    Node = {node, Value, Op},
    {Id, {tape, gleam@dict:insert(erlang:element(2, Tape), Id, Node), Id + 1}}.

-file("src/viva_math/autodiff_reverse.gleam", 99).
?DOC(" Register a new input variable on the tape.\n").
-spec input(tape(), float()) -> {integer(), tape()}.
input(Tape, Value) ->
    push(Tape, Value, input).

-file("src/viva_math/autodiff_reverse.gleam", 104).
?DOC(" Read the forward value of a node.\n").
-spec value(tape(), integer()) -> float().
value(Tape, Id) ->
    case gleam_stdlib:map_get(erlang:element(2, Tape), Id) of
        {ok, N} ->
            erlang:element(2, N);

        {error, _} ->
            +0.0
    end.

-file("src/viva_math/autodiff_reverse.gleam", 115).
-spec add(tape(), integer(), integer()) -> {integer(), tape()}.
add(Tape, A, B) ->
    push(Tape, value(Tape, A) + value(Tape, B), {add, A, B}).

-file("src/viva_math/autodiff_reverse.gleam", 119).
-spec sub(tape(), integer(), integer()) -> {integer(), tape()}.
sub(Tape, A, B) ->
    push(Tape, value(Tape, A) - value(Tape, B), {sub, A, B}).

-file("src/viva_math/autodiff_reverse.gleam", 123).
-spec mul(tape(), integer(), integer()) -> {integer(), tape()}.
mul(Tape, A, B) ->
    push(Tape, value(Tape, A) * value(Tape, B), {mul, A, B}).

-file("src/viva_math/autodiff_reverse.gleam", 127).
-spec 'div'(tape(), integer(), integer()) -> {integer(), tape()}.
'div'(Tape, A, B) ->
    push(Tape, case value(Tape, B) of
            +0.0 -> +0.0;
            -0.0 -> -0.0;
            Gleam@denominator -> value(Tape, A) / Gleam@denominator
        end, {'div', A, B}).

-file("src/viva_math/autodiff_reverse.gleam", 131).
-spec neg(tape(), integer()) -> {integer(), tape()}.
neg(Tape, A) ->
    push(Tape, +0.0 - value(Tape, A), {neg, A}).

-file("src/viva_math/autodiff_reverse.gleam", 135).
-spec scale(tape(), integer(), float()) -> {integer(), tape()}.
scale(Tape, A, S) ->
    push(Tape, value(Tape, A) * S, {scale, A, S}).

-file("src/viva_math/autodiff_reverse.gleam", 139).
-spec exp(tape(), integer()) -> {integer(), tape()}.
exp(Tape, A) ->
    push(Tape, math:exp(value(Tape, A)), {exp, A}).

-file("src/viva_math/autodiff_reverse.gleam", 143).
-spec ln(tape(), integer()) -> {integer(), tape()}.
ln(Tape, A) ->
    push(Tape, math:log(value(Tape, A)), {ln, A}).

-file("src/viva_math/autodiff_reverse.gleam", 147).
-spec sin(tape(), integer()) -> {integer(), tape()}.
sin(Tape, A) ->
    push(Tape, math:sin(value(Tape, A)), {sin, A}).

-file("src/viva_math/autodiff_reverse.gleam", 151).
-spec cos(tape(), integer()) -> {integer(), tape()}.
cos(Tape, A) ->
    push(Tape, math:cos(value(Tape, A)), {cos, A}).

-file("src/viva_math/autodiff_reverse.gleam", 155).
-spec tanh(tape(), integer()) -> {integer(), tape()}.
tanh(Tape, A) ->
    push(Tape, math:tanh(value(Tape, A)), {tanh, A}).

-file("src/viva_math/autodiff_reverse.gleam", 159).
-spec sigmoid(tape(), integer()) -> {integer(), tape()}.
sigmoid(Tape, A) ->
    V = value(Tape, A),
    S = case (1.0 + math:exp(+0.0 - V)) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> 1.0 / Gleam@denominator
    end,
    push(Tape, S, {sigmoid, A}).

-file("src/viva_math/autodiff_reverse.gleam", 165).
-spec pow(tape(), integer(), float()) -> {integer(), tape()}.
pow(Tape, A, N) ->
    push(Tape, math:pow(value(Tape, A), N), {pow, A, N}).

-file("src/viva_math/autodiff_reverse.gleam", 263).
-spec bump(gleam@dict:dict(integer(), float()), integer(), float()) -> gleam@dict:dict(integer(), float()).
bump(Grads, Id, Delta) ->
    Current = case gleam_stdlib:map_get(Grads, Id) of
        {ok, V} ->
            V;

        {error, _} ->
            +0.0
    end,
    gleam@dict:insert(Grads, Id, Current + Delta).

-file("src/viva_math/autodiff_reverse.gleam", 209).
-spec propagate(tape(), node_(), float(), gleam@dict:dict(integer(), float())) -> gleam@dict:dict(integer(), float()).
propagate(Tape, Node, G, Grads) ->
    case erlang:element(3, Node) of
        input ->
            Grads;

        {add, A, B} ->
            _pipe = Grads,
            _pipe@1 = bump(_pipe, A, G),
            bump(_pipe@1, B, G);

        {sub, A@1, B@1} ->
            _pipe@2 = Grads,
            _pipe@3 = bump(_pipe@2, A@1, G),
            bump(_pipe@3, B@1, +0.0 - G);

        {mul, A@2, B@2} ->
            Va = value(Tape, A@2),
            Vb = value(Tape, B@2),
            _pipe@4 = Grads,
            _pipe@5 = bump(_pipe@4, A@2, G * Vb),
            bump(_pipe@5, B@2, G * Va);

        {'div', A@3, B@3} ->
            Va@1 = value(Tape, A@3),
            Vb@1 = value(Tape, B@3),
            _pipe@6 = Grads,
            _pipe@7 = bump(_pipe@6, A@3, case Vb@1 of
                    +0.0 -> +0.0;
                    -0.0 -> -0.0;
                    Gleam@denominator -> G / Gleam@denominator
                end),
            bump(_pipe@7, B@3, +0.0 - (case (Vb@1 * Vb@1) of
                    +0.0 -> +0.0;
                    -0.0 -> -0.0;
                    Gleam@denominator@1 -> G * Va@1 / Gleam@denominator@1
                end));

        {neg, A@4} ->
            bump(Grads, A@4, +0.0 - G);

        {scale, A@5, S} ->
            bump(Grads, A@5, G * S);

        {exp, A@6} ->
            bump(Grads, A@6, G * math:exp(value(Tape, A@6)));

        {ln, A@7} ->
            bump(Grads, A@7, case value(Tape, A@7) of
                    +0.0 -> +0.0;
                    -0.0 -> -0.0;
                    Gleam@denominator@2 -> G / Gleam@denominator@2
                end);

        {sin, A@8} ->
            bump(Grads, A@8, G * math:cos(value(Tape, A@8)));

        {cos, A@9} ->
            bump(Grads, A@9, +0.0 - (G * math:sin(value(Tape, A@9))));

        {tanh, A@10} ->
            T = math:tanh(value(Tape, A@10)),
            bump(Grads, A@10, G * (1.0 - (T * T)));

        {sigmoid, A@11} ->
            V = value(Tape, A@11),
            S@1 = case (1.0 + math:exp(+0.0 - V)) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator@3 -> 1.0 / Gleam@denominator@3
            end,
            bump(Grads, A@11, (G * S@1) * (1.0 - S@1));

        {pow, A@12, N} ->
            Va@2 = value(Tape, A@12),
            bump(Grads, A@12, (G * N) * math:pow(Va@2, N - 1.0))
    end.

-file("src/viva_math/autodiff_reverse.gleam", 186).
-spec walk_back(tape(), integer(), gleam@dict:dict(integer(), float())) -> gleam@dict:dict(integer(), float()).
walk_back(Tape, Id, Grads) ->
    case Id < 0 of
        true ->
            Grads;

        false ->
            case gleam_stdlib:map_get(erlang:element(2, Tape), Id) of
                {error, _} ->
                    walk_back(Tape, Id - 1, Grads);

                {ok, Node} ->
                    G = case gleam_stdlib:map_get(Grads, Id) of
                        {ok, V} ->
                            V;

                        {error, _} ->
                            +0.0
                    end,
                    Grads_after = propagate(Tape, Node, G, Grads),
                    walk_back(Tape, Id - 1, Grads_after)
            end
    end.

-file("src/viva_math/autodiff_reverse.gleam", 177).
?DOC(
    " Compute ∂output/∂node for every node, given a scalar output node.\n"
    "\n"
    " Returns a Dict mapping each NodeId to its accumulated gradient. The\n"
    " `output` node itself has gradient 1.0.\n"
).
-spec backward(tape(), integer()) -> gleam@dict:dict(integer(), float()).
backward(Tape, Output) ->
    Grads = gleam@dict:insert(maps:new(), Output, 1.0),
    walk_back(Tape, Output, Grads).

-file("src/viva_math/autodiff_reverse.gleam", 276).
?DOC(" Extract gradient ∂output/∂input from the backward result.\n").
-spec grad_of(gleam@dict:dict(integer(), float()), integer()) -> float().
grad_of(Grads, Input) ->
    case gleam_stdlib:map_get(Grads, Input) of
        {ok, V} ->
            V;

        {error, _} ->
            +0.0
    end.

-file("src/viva_math/autodiff_reverse.gleam", 298).
-spec register_inputs(list(float()), tape(), list(integer())) -> {list(integer()),
    tape()}.
register_inputs(Values, Tape, Acc) ->
    case Values of
        [] ->
            {lists:reverse(Acc), Tape};

        [V | Rest] ->
            {Id, T} = input(Tape, V),
            register_inputs(Rest, T, [Id | Acc])
    end.

-file("src/viva_math/autodiff_reverse.gleam", 287).
?DOC(
    " Convenience: gradient of a scalar function with respect to many inputs.\n"
    "\n"
    " Runs the function on a fresh tape, performs the backward pass, and\n"
    " returns the gradient list aligned with the input list.\n"
).
-spec gradients(
    list(float()),
    fun((tape(), list(integer())) -> {integer(), tape()})
) -> list(float()).
gradients(Inputs, Build) ->
    Tape0 = empty_tape(),
    {Ids, Tape_with_inputs} = register_inputs(Inputs, Tape0, []),
    {Out, Final_tape} = Build(Tape_with_inputs, Ids),
    Grads = backward(Final_tape, Out),
    gleam@list:map(Ids, fun(Id) -> grad_of(Grads, Id) end).