-module(viva_math@ode).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_math/ode.gleam").
-export([euler/4, rk2_midpoint/4, rk2_heun/4, rk4/4, euler_maruyama/6, milstein/7, dop54/4, rkf45/4, velocity_verlet/4, leapfrog/4, position_verlet/4, yoshida4/4, integrate_symplectic/7, integrate/6, integrate_sde/7]).

-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(
    " Numerical ODE / SDE solvers.\n"
    "\n"
    " Single-step integrators for scalar systems `dx/dt = f(t, x)` and\n"
    " stochastic systems `dx = f(t, x) dt + g(t, x) dW`. Step functions are\n"
    " pure: they take the current `(t, x)` and produce the next state.\n"
    "\n"
    " ## Methods\n"
    "\n"
    " | Function          | Order  | Use case                                    |\n"
    " | ----------------- | ------ | ------------------------------------------- |\n"
    " | `euler`           | 1      | Quick, stiff systems with small dt          |\n"
    " | `rk2_midpoint`    | 2      | Mid-accuracy, half the cost of RK4          |\n"
    " | `rk2_heun`        | 2      | Trapezoidal/improved Euler                  |\n"
    " | `rk4`             | 4      | Workhorse for non-stiff systems             |\n"
    " | `euler_maruyama`  | 0.5*   | SDE with additive noise (* strong order)    |\n"
    " | `milstein`        | 1.0*   | SDE with multiplicative noise               |\n"
    "\n"
    " For full trajectories use `integrate` / `integrate_sde`.\n"
).

-file("src/viva_math/ode.gleam", 37).
?DOC(" Euler step: xₙ₊₁ = xₙ + dt · f(tₙ, xₙ).\n").
-spec euler(fun((float(), float()) -> float()), float(), float(), float()) -> float().
euler(F, T, X, Dt) ->
    X + (Dt * F(T, X)).

-file("src/viva_math/ode.gleam", 42).
?DOC(" Midpoint Runge-Kutta (RK2): evaluates f at the half-step midpoint.\n").
-spec rk2_midpoint(
    fun((float(), float()) -> float()),
    float(),
    float(),
    float()
) -> float().
rk2_midpoint(F, T, X, Dt) ->
    K1 = F(T, X),
    K2 = F(T + (Dt / 2.0), X + ((Dt * K1) / 2.0)),
    X + (Dt * K2).

-file("src/viva_math/ode.gleam", 49).
?DOC(" Heun (improved Euler) RK2: trapezoidal predictor-corrector.\n").
-spec rk2_heun(fun((float(), float()) -> float()), float(), float(), float()) -> float().
rk2_heun(F, T, X, Dt) ->
    K1 = F(T, X),
    K2 = F(T + Dt, X + (Dt * K1)),
    X + ((Dt * (K1 + K2)) / 2.0).

-file("src/viva_math/ode.gleam", 56).
?DOC(" Classical 4th-order Runge-Kutta.\n").
-spec rk4(fun((float(), float()) -> float()), float(), float(), float()) -> float().
rk4(F, T, X, Dt) ->
    K1 = F(T, X),
    K2 = F(T + (Dt / 2.0), X + ((Dt * K1) / 2.0)),
    K3 = F(T + (Dt / 2.0), X + ((Dt * K2) / 2.0)),
    K4 = F(T + Dt, X + (Dt * K3)),
    X + ((Dt * (((K1 + (2.0 * K2)) + (2.0 * K3)) + K4)) / 6.0).

-file("src/viva_math/ode.gleam", 71).
?DOC(
    " Euler-Maruyama step for dx = f dt + g dW with normal increment dW = √dt·Z.\n"
    "\n"
    " Returns the new state and the advanced PRNG seed.\n"
).
-spec euler_maruyama(
    fun((float(), float()) -> float()),
    fun((float(), float()) -> float()),
    float(),
    float(),
    float(),
    viva_math@random:seed()
) -> {float(), viva_math@random:seed()}.
euler_maruyama(F, G, T, X, Dt, Seed) ->
    Dt_safe = case Dt < +0.0 of
        true ->
            +0.0;

        false ->
            Dt
    end,
    {Z, S} = viva_math@random:standard_normal(Seed),
    Dw = Z * math:sqrt(Dt_safe),
    {(X + (F(T, X) * Dt_safe)) + (G(T, X) * Dw), S}.

-file("src/viva_math/ode.gleam", 95).
?DOC(
    " Milstein step — corrects Euler-Maruyama with the Itô derivative of g.\n"
    "\n"
    " dx = f dt + g dW + ½g·g'(t,x)·(dW² - dt)\n"
    "\n"
    " `dg_dx` is the partial derivative ∂g/∂x. If `g` doesn't depend on x\n"
    " (additive noise), use `euler_maruyama` instead — they coincide.\n"
).
-spec milstein(
    fun((float(), float()) -> float()),
    fun((float(), float()) -> float()),
    fun((float(), float()) -> float()),
    float(),
    float(),
    float(),
    viva_math@random:seed()
) -> {float(), viva_math@random:seed()}.
milstein(F, G, Dg_dx, T, X, Dt, Seed) ->
    {Z, S} = viva_math@random:standard_normal(Seed),
    Dw = Z * math:sqrt(Dt),
    G_val = G(T, X),
    Correction = ((0.5 * G_val) * Dg_dx(T, X)) * ((Dw * Dw) - Dt),
    {((X + (F(T, X) * Dt)) + (G_val * Dw)) + Correction, S}.

-file("src/viva_math/ode.gleam", 139).
?DOC(
    " One step of Dormand-Prince 5(4) with embedded error estimate.\n"
    "\n"
    " Returns `#(x_new, err)` where `x_new` is the 5th-order solution and\n"
    " `err` is the absolute difference from the embedded 4th-order estimate,\n"
    " suitable for step-size control. Per-step truncation error is O(dt⁶).\n"
).
-spec dop54(fun((float(), float()) -> float()), float(), float(), float()) -> {float(),
    float()}.
dop54(F, T, X, Dt) ->
    K1 = F(T, X),
    K2 = F(T + (Dt / 5.0), X + ((Dt * K1) / 5.0)),
    K3 = F(
        T + ((3.0 * Dt) / 10.0),
        X + (Dt * (((3.0 * K1) / 40.0) + ((9.0 * K2) / 40.0)))
    ),
    K4 = F(
        T + ((4.0 * Dt) / 5.0),
        X + (Dt * ((((44.0 * K1) / 45.0) - ((56.0 * K2) / 15.0)) + ((32.0 * K3)
        / 9.0)))
    ),
    K5 = F(
        T + ((8.0 * Dt) / 9.0),
        X + (Dt * (((((19372.0 * K1) / 6561.0) - ((25360.0 * K2) / 2187.0)) + ((64448.0
        * K3)
        / 6561.0))
        - ((212.0 * K4) / 729.0)))
    ),
    K6 = F(
        T + Dt,
        X + (Dt * ((((((9017.0 * K1) / 3168.0) - ((355.0 * K2) / 33.0)) + ((46732.0
        * K3)
        / 5247.0))
        + ((49.0 * K4) / 176.0))
        - ((5103.0 * K5) / 18656.0)))
    ),
    X_new = X + (Dt * ((((((35.0 * K1) / 384.0) + ((500.0 * K3) / 1113.0)) + ((125.0
    * K4)
    / 192.0))
    - ((2187.0 * K5) / 6784.0))
    + ((11.0 * K6) / 84.0))),
    K7 = F(T + Dt, X_new),
    X_alt = X + (Dt * (((((((5179.0 * K1) / 57600.0) + ((7571.0 * K3) / 16695.0))
    + ((393.0 * K4) / 640.0))
    - ((92097.0 * K5) / 339200.0))
    + ((187.0 * K6) / 2100.0))
    + (K7 / 40.0))),
    Err = case X_new >= X_alt of
        true ->
            X_new - X_alt;

        false ->
            X_alt - X_new
    end,
    {X_new, Err}.

-file("src/viva_math/ode.gleam", 261).
?DOC(
    " Deprecated alias retained for backwards compatibility. Returns the same\n"
    " pair as `dop54`. To be removed once external callers migrate.\n"
    " One step of RKF45 with an error estimate.\n"
    "\n"
    " Returns `#(x5, error)` where `x5` is the 5th-order estimate and `error` is\n"
    " |x5 - x4|, useful for step-size control.\n"
).
-spec rkf45(fun((float(), float()) -> float()), float(), float(), float()) -> {float(),
    float()}.
rkf45(F, T, X, Dt) ->
    K1 = Dt * F(T, X),
    K2 = Dt * F(T + (Dt / 4.0), X + (K1 / 4.0)),
    K3 = Dt * F(
        T + ((3.0 * Dt) / 8.0),
        (X + ((3.0 * K1) / 32.0)) + ((9.0 * K2) / 32.0)
    ),
    K4 = Dt * F(
        T + ((12.0 * Dt) / 13.0),
        ((X + ((1932.0 * K1) / 2197.0)) - ((7200.0 * K2) / 2197.0)) + ((7296.0 * K3)
        / 2197.0)
    ),
    K5 = Dt * F(
        T + Dt,
        (((X + ((439.0 * K1) / 216.0)) - (8.0 * K2)) + ((3680.0 * K3) / 513.0))
        - ((845.0 * K4) / 4104.0)
    ),
    K6 = Dt * F(
        T + (Dt / 2.0),
        ((((X - ((8.0 * K1) / 27.0)) + (2.0 * K2)) - ((3544.0 * K3) / 2565.0)) + ((1859.0
        * K4)
        / 4104.0))
        - ((11.0 * K5) / 40.0)
    ),
    X4 = (((X + ((25.0 * K1) / 216.0)) + ((1408.0 * K3) / 2565.0)) + ((2197.0 * K4)
    / 4104.0))
    - (K5 / 5.0),
    X5 = ((((X + ((16.0 * K1) / 135.0)) + ((6656.0 * K3) / 12825.0)) + ((28561.0
    * K4)
    / 56430.0))
    - ((9.0 * K5) / 50.0))
    + ((2.0 * K6) / 55.0),
    Err = case X5 >= X4 of
        true ->
            X5 - X4;

        false ->
            X4 - X5
    end,
    {X5, Err}.

-file("src/viva_math/ode.gleam", 372).
?DOC(" Velocity Verlet — order 2, symplectic, time-reversible.\n").
-spec velocity_verlet(fun((float()) -> float()), float(), float(), float()) -> {float(),
    float()}.
velocity_verlet(Force, Q, V, Dt) ->
    A = Force(Q),
    Q_new = (Q + (V * Dt)) + (((0.5 * A) * Dt) * Dt),
    A_new = Force(Q_new),
    V_new = V + ((0.5 * (A + A_new)) * Dt),
    {Q_new, V_new}.

-file("src/viva_math/ode.gleam", 386).
?DOC(" Leapfrog (kick-drift-kick) — order 2, symplectic.\n").
-spec leapfrog(fun((float()) -> float()), float(), float(), float()) -> {float(),
    float()}.
leapfrog(Force, Q, V, Dt) ->
    V_half = V + ((0.5 * Force(Q)) * Dt),
    Q_new = Q + (V_half * Dt),
    V_new = V_half + ((0.5 * Force(Q_new)) * Dt),
    {Q_new, V_new}.

-file("src/viva_math/ode.gleam", 399).
?DOC(" Position Verlet — `q(t+dt) = 2·q(t) - q(t-dt) + a(q(t))·dt²`.\n").
-spec position_verlet(fun((float()) -> float()), float(), float(), float()) -> float().
position_verlet(Force, Q_prev, Q, Dt) ->
    ((2.0 * Q) - Q_prev) + ((Force(Q) * Dt) * Dt).

-file("src/viva_math/ode.gleam", 413).
?DOC(
    " Yoshida 4th-order symplectic integrator.\n"
    "\n"
    " Composes three leapfrog steps with Yoshida (1990) coefficients. Order 4\n"
    " global error, ~3× the cost of leapfrog per step but allows ~10× larger\n"
    " `dt` for the same accuracy.\n"
).
-spec yoshida4(fun((float()) -> float()), float(), float(), float()) -> {float(),
    float()}.
yoshida4(Force, Q, V, Dt) ->
    Cbrt_2 = 1.2599210498948732,
    W1 = case (2.0 - Cbrt_2) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> 1.0 / Gleam@denominator
    end,
    W0 = +0.0 - (Cbrt_2 * W1),
    {Q1, V1} = leapfrog(Force, Q, V, W1 * Dt),
    {Q2, V2} = leapfrog(Force, Q1, V1, W0 * Dt),
    leapfrog(Force, Q2, V2, W1 * Dt).

-file("src/viva_math/ode.gleam", 440).
-spec symplectic_loop(
    fun((fun((float()) -> float()), float(), float(), float()) -> {float(),
        float()}),
    fun((float()) -> float()),
    float(),
    float(),
    float(),
    float(),
    integer(),
    list({float(), float(), float()})
) -> list({float(), float(), float()}).
symplectic_loop(Method, Force, T, Q, V, Dt, Steps, Acc) ->
    case Steps =< 0 of
        true ->
            lists:reverse(Acc);

        false ->
            {Q_new, V_new} = Method(Force, Q, V, Dt),
            T_new = T + Dt,
            symplectic_loop(
                Method,
                Force,
                T_new,
                Q_new,
                V_new,
                Dt,
                Steps - 1,
                [{T_new, Q_new, V_new} | Acc]
            )
    end.

-file("src/viva_math/ode.gleam", 428).
?DOC(" Integrate a Hamiltonian system with a symplectic method.\n").
-spec integrate_symplectic(
    fun((fun((float()) -> float()), float(), float(), float()) -> {float(),
        float()}),
    fun((float()) -> float()),
    float(),
    float(),
    float(),
    float(),
    integer()
) -> list({float(), float(), float()}).
integrate_symplectic(Method, Force, T0, Q0, V0, Dt, Steps) ->
    symplectic_loop(Method, Force, T0, Q0, V0, Dt, Steps, [{T0, Q0, V0}]).

-file("src/viva_math/ode.gleam", 479).
-spec integrate_loop(
    fun((fun((float(), float()) -> float()), float(), float(), float()) -> float()),
    fun((float(), float()) -> float()),
    float(),
    float(),
    float(),
    integer(),
    list({float(), float()})
) -> list({float(), float()}).
integrate_loop(Method, F, T, X, Dt, Steps, Acc) ->
    case Steps =< 0 of
        true ->
            lists:reverse(Acc);

        false ->
            Next_t = T + Dt,
            Next_x = Method(F, T, X, Dt),
            integrate_loop(
                Method,
                F,
                Next_t,
                Next_x,
                Dt,
                Steps - 1,
                [{Next_t, Next_x} | Acc]
            )
    end.

-file("src/viva_math/ode.gleam", 468).
?DOC(" Integrate a deterministic ODE with a fixed-step method.\n").
-spec integrate(
    fun((fun((float(), float()) -> float()), float(), float(), float()) -> float()),
    fun((float(), float()) -> float()),
    float(),
    float(),
    float(),
    integer()
) -> list({float(), float()}).
integrate(Method, F, T0, X0, Dt, Steps) ->
    integrate_loop(Method, F, T0, X0, Dt, Steps, [{T0, X0}]).

-file("src/viva_math/ode.gleam", 514).
-spec integrate_sde_loop(
    fun((float(), float()) -> float()),
    fun((float(), float()) -> float()),
    float(),
    float(),
    float(),
    integer(),
    viva_math@random:seed(),
    list({float(), float()})
) -> {list({float(), float()}), viva_math@random:seed()}.
integrate_sde_loop(F, G, T, X, Dt, Steps, Seed, Acc) ->
    case Steps =< 0 of
        true ->
            {lists:reverse(Acc), Seed};

        false ->
            {Next_x, Next_seed} = euler_maruyama(F, G, T, X, Dt, Seed),
            Next_t = T + Dt,
            integrate_sde_loop(
                F,
                G,
                Next_t,
                Next_x,
                Dt,
                Steps - 1,
                Next_seed,
                [{Next_t, Next_x} | Acc]
            )
    end.

-file("src/viva_math/ode.gleam", 502).
?DOC(" Integrate a stochastic system with Euler-Maruyama.\n").
-spec integrate_sde(
    fun((float(), float()) -> float()),
    fun((float(), float()) -> float()),
    float(),
    float(),
    float(),
    integer(),
    viva_math@random:seed()
) -> {list({float(), float()}), viva_math@random:seed()}.
integrate_sde(F, G, T0, X0, Dt, Steps, Seed) ->
    integrate_sde_loop(F, G, T0, X0, Dt, Steps, Seed, [{T0, X0}]).