viva_math

-module(viva_math@autodiff).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_math/autodiff.gleam").
-export([var/1, constant/1, lift1/3, add/2, sub/2, mul/2, 'div'/2, neg/1, scale/2, add_scalar/2, exp/1, ln/1, sqrt/1, pow/2, sin/1, cos/1, tanh/1, sigmoid/1, gelu/1, relu/1, grad/2, value_and_grad/2, const3/1, var3_x/1, var3_y/1, var3_z/1, add3/2, mul3/2, exp3/1, gradient3/4, jacobian/2]).
-export_type([dual/0, dual3/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(
    " Forward-mode automatic differentiation via dual numbers.\n"
    "\n"
    " A dual number `a + b·ε` (where `ε² = 0`) carries both a value and an\n"
    " infinitesimal derivative. Applying `f` to a `Dual(x, 1)` yields\n"
    " `Dual(f(x), f'(x))` — **exact** gradients without symbolic\n"
    " manipulation or finite-difference truncation error.\n"
    "\n"
    " ## When to use\n"
    "\n"
    " - You need an exact gradient at a point.\n"
    " - The function is composed from arithmetic ops + `exp/log/sin/cos/...`.\n"
    " - Cost of evaluating `f` once with duals ≈ 2-3× the regular cost.\n"
    "\n"
    " For Vec3 gradients (gradient of a scalar field over PAD space), use\n"
    " `Dual3` which carries three partials in parallel.\n"
    "\n"
    " ## Example\n"
    "\n"
    " ```gleam\n"
    " import viva_math/autodiff as ad\n"
    "\n"
    " // f(x) = sin(x²)\n"
    " let x = ad.var(2.0)\n"
    " let result = ad.sin(ad.mul(x, x))\n"
    " // result.value ≈ sin(4) ≈ -0.756\n"
    " // result.tangent = 2x · cos(x²) ≈ -2.614\n"
    " ```\n"
    "\n"
    " ## References\n"
    "\n"
    " - Pearlmutter & Siskind (2008) \"Reverse-Mode AD in a Functional Framework\"\n"
    " - Wengert (1964) \"A simple automatic derivative evaluation program\"\n"
    " - JAX `jax.jvp`, Zygote.jl, Stan's Math library\n"
).

-type dual() :: {dual, float(), float()}.

-type dual3() :: {dual3, float(), float(), float(), float()}.

-file("src/viva_math/autodiff.gleam", 52).
?DOC(
    " A dual number representing an independent variable.\n"
    " `var(x)` ↔ `x + 1·ε` so that ∂x/∂x = 1.\n"
).
-spec var(float()) -> dual().
var(Value) ->
    {dual, Value, 1.0}.

-file("src/viva_math/autodiff.gleam", 57).
?DOC(" Constant dual: tangent = 0.\n").
-spec constant(float()) -> dual().
constant(Value) ->
    {dual, Value, +0.0}.

-file("src/viva_math/autodiff.gleam", 62).
?DOC(" Lift any unary function `f` together with its derivative `f'` to a dual.\n").
-spec lift1(dual(), fun((float()) -> float()), fun((float()) -> float())) -> dual().
lift1(D, F, Df) ->
    {dual,
        F(erlang:element(2, D)),
        Df(erlang:element(2, D)) * erlang:element(3, D)}.

-file("src/viva_math/autodiff.gleam", 68).
-spec add(dual(), dual()) -> dual().
add(A, B) ->
    {dual,
        erlang:element(2, A) + erlang:element(2, B),
        erlang:element(3, A) + erlang:element(3, B)}.

-file("src/viva_math/autodiff.gleam", 72).
-spec sub(dual(), dual()) -> dual().
sub(A, B) ->
    {dual,
        erlang:element(2, A) - erlang:element(2, B),
        erlang:element(3, A) - erlang:element(3, B)}.

-file("src/viva_math/autodiff.gleam", 76).
-spec mul(dual(), dual()) -> dual().
mul(A, B) ->
    {dual,
        erlang:element(2, A) * erlang:element(2, B),
        (erlang:element(3, A) * erlang:element(2, B)) + (erlang:element(2, A) * erlang:element(
            3,
            B
        ))}.

-file("src/viva_math/autodiff.gleam", 81).
-spec 'div'(dual(), dual()) -> dual().
'div'(A, B) ->
    Denom = erlang:element(2, B) * erlang:element(2, B),
    {dual, case erlang:element(2, B) of
            +0.0 -> +0.0;
            -0.0 -> -0.0;
            Gleam@denominator -> erlang:element(2, A) / Gleam@denominator
        end, case Denom of
            +0.0 -> +0.0;
            -0.0 -> -0.0;
            Gleam@denominator@1 -> ((erlang:element(3, A) * erlang:element(2, B))
            - (erlang:element(2, A) * erlang:element(3, B)))
            / Gleam@denominator@1
        end}.

-file("src/viva_math/autodiff.gleam", 90).
-spec neg(dual()) -> dual().
neg(A) ->
    {dual, +0.0 - erlang:element(2, A), +0.0 - erlang:element(3, A)}.

-file("src/viva_math/autodiff.gleam", 94).
-spec scale(dual(), float()) -> dual().
scale(A, S) ->
    {dual, erlang:element(2, A) * S, erlang:element(3, A) * S}.

-file("src/viva_math/autodiff.gleam", 98).
-spec add_scalar(dual(), float()) -> dual().
add_scalar(A, S) ->
    {dual, erlang:element(2, A) + S, erlang:element(3, A)}.

-file("src/viva_math/autodiff.gleam", 104).
-spec exp(dual()) -> dual().
exp(A) ->
    V = math:exp(erlang:element(2, A)),
    {dual, V, V * erlang:element(3, A)}.

-file("src/viva_math/autodiff.gleam", 110).
-spec ln(dual()) -> dual().
ln(A) ->
    {dual, math:log(erlang:element(2, A)), case erlang:element(2, A) of
            +0.0 -> +0.0;
            -0.0 -> -0.0;
            Gleam@denominator -> erlang:element(3, A) / Gleam@denominator
        end}.

-file("src/viva_math/autodiff.gleam", 114).
-spec sqrt(dual()) -> dual().
sqrt(A) ->
    V = math:sqrt(erlang:element(2, A)),
    {dual, V, case (2.0 * V) of
            +0.0 -> +0.0;
            -0.0 -> -0.0;
            Gleam@denominator -> erlang:element(3, A) / Gleam@denominator
        end}.

-file("src/viva_math/autodiff.gleam", 119).
-spec pow(dual(), float()) -> dual().
pow(A, N) ->
    V = math:pow(erlang:element(2, A), N),
    {dual,
        V,
        (N * math:pow(erlang:element(2, A), N - 1.0)) * erlang:element(3, A)}.

-file("src/viva_math/autodiff.gleam", 125).
-spec sin(dual()) -> dual().
sin(A) ->
    {dual,
        math:sin(erlang:element(2, A)),
        math:cos(erlang:element(2, A)) * erlang:element(3, A)}.

-file("src/viva_math/autodiff.gleam", 129).
-spec cos(dual()) -> dual().
cos(A) ->
    {dual,
        math:cos(erlang:element(2, A)),
        +0.0 - (math:sin(erlang:element(2, A)) * erlang:element(3, A))}.

-file("src/viva_math/autodiff.gleam", 133).
-spec tanh(dual()) -> dual().
tanh(A) ->
    T = math:tanh(erlang:element(2, A)),
    {dual, T, (1.0 - (T * T)) * erlang:element(3, A)}.

-file("src/viva_math/autodiff.gleam", 138).
-spec sigmoid(dual()) -> dual().
sigmoid(A) ->
    S = viva_math@scalar:sigmoid(erlang:element(2, A)),
    {dual, S, (S * (1.0 - S)) * erlang:element(3, A)}.

-file("src/viva_math/autodiff.gleam", 143).
-spec gelu(dual()) -> dual().
gelu(A) ->
    Phi = 0.5 * (1.0 + math:erf(erlang:element(2, A) * 0.7071067811865475)),
    Phi_pdf = 0.3989422804014327 * math:exp(
        (-0.5 * erlang:element(2, A)) * erlang:element(2, A)
    ),
    G = erlang:element(2, A) * Phi,
    {dual, G, (Phi + (erlang:element(2, A) * Phi_pdf)) * erlang:element(3, A)}.

-file("src/viva_math/autodiff.gleam", 152).
-spec relu(dual()) -> dual().
relu(A) ->
    case erlang:element(2, A) > +0.0 of
        true ->
            A;

        false ->
            {dual, +0.0, +0.0}
    end.

-file("src/viva_math/autodiff.gleam", 164).
?DOC(" Compute f'(x) at a point by lifting `x` to a dual.\n").
-spec grad(fun((dual()) -> dual()), float()) -> float().
grad(F, X) ->
    erlang:element(3, F(var(X))).

-file("src/viva_math/autodiff.gleam", 169).
?DOC(" Compute both f(x) and f'(x) in a single pass.\n").
-spec value_and_grad(fun((dual()) -> dual()), float()) -> {float(), float()}.
value_and_grad(F, X) ->
    R = F(var(X)),
    {erlang:element(2, R), erlang:element(3, R)}.

-file("src/viva_math/autodiff.gleam", 185).
?DOC(" Treat `value` as a constant: all partials = 0.\n").
-spec const3(float()) -> dual3().
const3(Value) ->
    {dual3, Value, +0.0, +0.0, +0.0}.

-file("src/viva_math/autodiff.gleam", 190).
?DOC(" The independent variable x: ∂x/∂x = 1, others = 0.\n").
-spec var3_x(float()) -> dual3().
var3_x(Value) ->
    {dual3, Value, 1.0, +0.0, +0.0}.

-file("src/viva_math/autodiff.gleam", 194).
-spec var3_y(float()) -> dual3().
var3_y(Value) ->
    {dual3, Value, +0.0, 1.0, +0.0}.

-file("src/viva_math/autodiff.gleam", 198).
-spec var3_z(float()) -> dual3().
var3_z(Value) ->
    {dual3, Value, +0.0, +0.0, 1.0}.

-file("src/viva_math/autodiff.gleam", 202).
-spec add3(dual3(), dual3()) -> dual3().
add3(A, B) ->
    {dual3,
        erlang:element(2, A) + erlang:element(2, B),
        erlang:element(3, A) + erlang:element(3, B),
        erlang:element(4, A) + erlang:element(4, B),
        erlang:element(5, A) + erlang:element(5, B)}.

-file("src/viva_math/autodiff.gleam", 211).
-spec mul3(dual3(), dual3()) -> dual3().
mul3(A, B) ->
    {dual3,
        erlang:element(2, A) * erlang:element(2, B),
        (erlang:element(3, A) * erlang:element(2, B)) + (erlang:element(2, A) * erlang:element(
            3,
            B
        )),
        (erlang:element(4, A) * erlang:element(2, B)) + (erlang:element(2, A) * erlang:element(
            4,
            B
        )),
        (erlang:element(5, A) * erlang:element(2, B)) + (erlang:element(2, A) * erlang:element(
            5,
            B
        ))}.

-file("src/viva_math/autodiff.gleam", 220).
-spec exp3(dual3()) -> dual3().
exp3(A) ->
    V = math:exp(erlang:element(2, A)),
    {dual3,
        V,
        V * erlang:element(3, A),
        V * erlang:element(4, A),
        V * erlang:element(5, A)}.

-file("src/viva_math/autodiff.gleam", 226).
?DOC(" Gradient ∇f at point (x, y, z).\n").
-spec gradient3(
    fun((dual3(), dual3(), dual3()) -> dual3()),
    float(),
    float(),
    float()
) -> {float(), float(), float()}.
gradient3(F, X, Y, Z) ->
    R = F(var3_x(X), var3_y(Y), var3_z(Z)),
    {erlang:element(3, R), erlang:element(4, R), erlang:element(5, R)}.

-file("src/viva_math/autodiff.gleam", 255).
-spec build_input(list(float()), integer(), integer(), list(dual())) -> list(dual()).
build_input(Point, Active, I, Acc) ->
    case Point of
        [] ->
            lists:reverse(Acc);

        [X | Rest] ->
            D = case I =:= Active of
                true ->
                    {dual, X, 1.0};

                false ->
                    {dual, X, +0.0}
            end,
            build_input(Rest, Active, I + 1, [D | Acc])
    end.

-file("src/viva_math/autodiff.gleam", 283).
-spec range_loop(integer(), integer(), list(integer())) -> list(integer()).
range_loop(From, To, Acc) ->
    case From > To of
        true ->
            lists:reverse(Acc);

        false ->
            range_loop(From + 1, To, [From | Acc])
    end.

-file("src/viva_math/autodiff.gleam", 279).
-spec range_int(integer(), integer()) -> list(integer()).
range_int(From, To) ->
    range_loop(From, To, []).

-file("src/viva_math/autodiff.gleam", 242).
?DOC(
    " Generic n-d Jacobian via column-by-column forward AD. Given a function\n"
    " `f: ℝⁿ → ℝᵐ` represented as `fn(Dual, ..., Dual) -> List(Dual)`, returns\n"
    " each row of the Jacobian by sweeping the unit tangent through each input.\n"
    "\n"
    " This is the canonical forward-mode strategy: O(n) evaluations of `f` for\n"
    " a full Jacobian, optimal when n ≤ m.\n"
).
-spec jacobian(fun((list(dual())) -> list(dual())), list(float())) -> list(list(float())).
jacobian(F, Point) ->
    N = erlang:length(Point),
    _pipe = range_int(0, N - 1),
    gleam@list:map(
        _pipe,
        fun(I) ->
            Inputs = build_input(Point, I, 0, []),
            Outputs = F(Inputs),
            gleam@list:map(Outputs, fun(D) -> erlang:element(3, D) end)
        end
    ).