viva_math

-module(viva_math@ou).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_math/ou.gleam").
-export([is_valid/1, is_valid_vec3/1, step/4, simulate/5, mean_at/3, variance_at/3, stationary_variance/1, stationary_std/1, autocovariance/2, half_life/1, step_vec3/4, simulate_vec3/5, mean_at_vec3/3, variance_at_vec3/3, stationary_variance_vec3/1]).
-export_type([o_u_params1_d/0, o_u_params_vec3/0, kernel/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(
    " Ornstein-Uhlenbeck mood dynamics.\n"
    "\n"
    " Mean-reverting stochastic process for affective dynamics — the canonical\n"
    " model for emotion regulation toward a baseline. Underlies VIVA's\n"
    " homeostatic emotional decay.\n"
    "\n"
    " **SDE**: `dX_t = θ(μ - X_t) dt + σ dW_t`\n"
    "\n"
    " Parameters:\n"
    " - `theta` (θ) — mean-reversion speed (> 0). Larger = faster return to μ.\n"
    " - `mu`    (μ) — long-run mean (the attractor).\n"
    " - `sigma` (σ) — diffusion (volatility, ≥ 0).\n"
    "\n"
    " ## Analytical properties\n"
    "\n"
    " Given `X_0 = x_0`:\n"
    "   - `E[X_t]   = μ + (x_0 − μ) · e^(−θt)`\n"
    "   - `Var[X_t] = σ² / (2θ) · (1 − e^(−2θt))`\n"
    "   - Stationary: `X_∞ ~ N(μ, σ²/(2θ))`\n"
    "   - Autocovariance at lag `τ`: `σ²/(2θ) · e^(−θ|τ|)`\n"
    "   - Half-life of expectation: `ln(2) / θ`\n"
    "\n"
    " ## Integration\n"
    "\n"
    " `step` uses the **exact transition kernel** (Doob 1942) — no\n"
    " discretization error regardless of `dt`. Closed form:\n"
    "\n"
    " ```\n"
    " X_{t+Δ} = μ + (X_t − μ)·e^(−θΔ) + σ·sqrt((1 − e^(−2θΔ))/(2θ)) · Z\n"
    " ```\n"
    "\n"
    " where `Z ~ N(0, 1)`. For Euler-Maruyama on the same SDE, use\n"
    " `ode.euler_maruyama` directly with a custom drift/diffusion.\n"
    "\n"
    " ## References\n"
    "\n"
    " - Uhlenbeck & Ornstein (1930) — *On the theory of Brownian motion*\n"
    " - Oravecz, Tuerlinckx & Vandekerckhove (2009) — *Ornstein-Uhlenbeck\n"
    "   Process in Affective Dynamics*\n"
    " - Doob (1942) — *The Brownian Movement and Stochastic Equations*\n"
).

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

-type o_u_params_vec3() :: {o_u_params_vec3,
        viva_math@vector:vec3(),
        viva_math@vector:vec3(),
        viva_math@vector:vec3()}.

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

-file("src/viva_math/ou.gleam", 73).
?DOC(
    " Check whether 1D parameters are physically meaningful.\n"
    "\n"
    " Requires `theta > 0` (otherwise no mean-reversion) and `sigma >= 0`.\n"
).
-spec is_valid(o_u_params1_d()) -> boolean().
is_valid(Params) ->
    (erlang:element(2, Params) > +0.0) andalso (erlang:element(4, Params) >= +0.0).

-file("src/viva_math/ou.gleam", 79).
?DOC(
    " Vec3 validity — every component of `theta` strictly positive and every\n"
    " component of `sigma` non-negative.\n"
).
-spec is_valid_vec3(o_u_params_vec3()) -> boolean().
is_valid_vec3(Params) ->
    (((((erlang:element(2, erlang:element(2, Params)) > +0.0) andalso (erlang:element(
        3,
        erlang:element(2, Params)
    )
    > +0.0))
    andalso (erlang:element(4, erlang:element(2, Params)) > +0.0))
    andalso (erlang:element(2, erlang:element(4, Params)) >= +0.0))
    andalso (erlang:element(3, erlang:element(4, Params)) >= +0.0))
    andalso (erlang:element(4, erlang:element(4, Params)) >= +0.0).

-file("src/viva_math/ou.gleam", 103).
?DOC(
    " One step of the exact OU transition kernel.\n"
    "\n"
    " `X_{t+dt} = μ + (X_t − μ)·e^(−θ·dt) + σ·sqrt((1 − e^(−2θ·dt))/(2θ)) · Z`\n"
    "\n"
    " `Z ~ N(0, 1)`. No discretization error: works correctly even for large `dt`.\n"
    "\n"
    " **Caller-validated inputs**: callers must ensure `theta > 0`, `sigma >= 0`,\n"
    " and `dt >= 0` (use `is_valid` for the params). With `dt < 0` the variance\n"
    " term goes negative and `std_term` silently collapses to `0.0`, producing\n"
    " a deterministic backward step that consumes a normal draw without using\n"
    " it — not a physically meaningful transition.\n"
).
-spec step(o_u_params1_d(), float(), float(), viva_math@random:seed()) -> {float(),
    viva_math@random:seed()}.
step(Params, X, Dt, Seed) ->
    {o_u_params1_d, Theta, Mu, Sigma} = Params,
    Decay = math:exp(+0.0 - (Theta * Dt)),
    Drift_term = Mu + ((X - Mu) * Decay),
    Var_term = case (2.0 * Theta) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> (Sigma * Sigma) * (+0.0 - viva_math@scalar:expm1(
            +0.0 - ((2.0 * Theta) * Dt)
        ))
        / Gleam@denominator
    end,
    Std_term = case gleam@float:square_root(Var_term) of
        {ok, S} ->
            S;

        {error, _} ->
            +0.0
    end,
    {Z, New_seed} = viva_math@random:standard_normal(Seed),
    {Drift_term + (Std_term * Z), New_seed}.

-file("src/viva_math/ou.gleam", 156).
-spec simulate_loop(
    float(),
    float(),
    float(),
    float(),
    integer(),
    viva_math@random:seed(),
    list(float())
) -> {list(float()), viva_math@random:seed()}.
simulate_loop(Decay, Mu, Std, X, N, Seed, Acc) ->
    case N =< 0 of
        true ->
            {Acc, Seed};

        false ->
            {Z, S_next} = viva_math@random:standard_normal(Seed),
            X_next = (Mu + ((X - Mu) * Decay)) + (Std * Z),
            simulate_loop(Decay, Mu, Std, X_next, N - 1, S_next, [X_next | Acc])
    end.

-file("src/viva_math/ou.gleam", 134).
?DOC(
    " Simulate `n` steps starting from `x0` with constant time-step `dt`.\n"
    "\n"
    " Returns the trajectory **excluding** the initial point (length `n`) and the\n"
    " final seed for chaining.\n"
    "\n"
    " Pre-computes the transition kernel (`decay`, `std`) once — the loop only\n"
    " does a multiply-add and a normal draw per step.\n"
).
-spec simulate(
    o_u_params1_d(),
    float(),
    float(),
    integer(),
    viva_math@random:seed()
) -> {list(float()), viva_math@random:seed()}.
simulate(Params, X0, Dt, N, Seed) ->
    {o_u_params1_d, Theta, Mu, Sigma} = Params,
    Decay = math:exp(+0.0 - (Theta * Dt)),
    Var_term = case (2.0 * Theta) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> (Sigma * Sigma) * (+0.0 - viva_math@scalar:expm1(
            +0.0 - ((2.0 * Theta) * Dt)
        ))
        / Gleam@denominator
    end,
    Std_term = case gleam@float:square_root(Var_term) of
        {ok, S} ->
            S;

        {error, _} ->
            +0.0
    end,
    {Traj, S@1} = simulate_loop(Decay, Mu, Std_term, X0, N, Seed, []),
    {lists:reverse(Traj), S@1}.

-file("src/viva_math/ou.gleam", 182).
?DOC(
    " Closed-form `E[X_t | X_0 = x0]`.\n"
    "\n"
    " `μ + (x0 − μ) · e^(−θ·t)`\n"
).
-spec mean_at(o_u_params1_d(), float(), float()) -> float().
mean_at(Params, X0, T) ->
    {o_u_params1_d, Theta, Mu, _} = Params,
    Decay = math:exp(+0.0 - (Theta * T)),
    Mu + ((X0 - Mu) * Decay).

-file("src/viva_math/ou.gleam", 199).
?DOC(
    " Closed-form `Var[X_t | X_0 = x0]`.\n"
    "\n"
    " `σ² / (2θ) · (1 − e^(−2θ·t))`\n"
    "\n"
    " **Note**: the conditional variance is **independent of `x0`** — OU's noise\n"
    " is additive Brownian, so all `x0`-dependence is absorbed into the mean.\n"
    " The parameter is kept in the signature only to mirror `mean_at` and\n"
    " `variance_at_vec3` for API symmetry; pass any value.\n"
    "\n"
    " Routed through `scalar.expm1` so the Brownian limit `σ²·t` (as `θ·t → 0`)\n"
    " is recovered without catastrophic cancellation.\n"
).
-spec variance_at(o_u_params1_d(), float(), float()) -> float().
variance_at(Params, _, T) ->
    {o_u_params1_d, Theta, _, Sigma} = Params,
    case (2.0 * Theta) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> (Sigma * Sigma) * (+0.0 - viva_math@scalar:expm1(
            +0.0 - ((2.0 * Theta) * T)
        ))
        / Gleam@denominator
    end.

-file("src/viva_math/ou.gleam", 209).
?DOC(" Stationary variance `σ² / (2θ)` — the variance of `X_∞ ~ N(μ, σ²/(2θ))`.\n").
-spec stationary_variance(o_u_params1_d()) -> float().
stationary_variance(Params) ->
    {o_u_params1_d, Theta, _, Sigma} = Params,
    case (2.0 * Theta) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> Sigma * Sigma / Gleam@denominator
    end.

-file("src/viva_math/ou.gleam", 215).
?DOC(" Stationary standard deviation.\n").
-spec stationary_std(o_u_params1_d()) -> float().
stationary_std(Params) ->
    case gleam@float:square_root(stationary_variance(Params)) of
        {ok, S} ->
            S;

        {error, _} ->
            +0.0
    end.

-file("src/viva_math/ou.gleam", 225).
?DOC(
    " Autocovariance at lag `τ` (in time units).\n"
    "\n"
    " `Cov(X_s, X_{s+τ}) = σ²/(2θ) · e^(−θ·|τ|)` (stationary regime).\n"
).
-spec autocovariance(o_u_params1_d(), float()) -> float().
autocovariance(Params, Lag) ->
    Abs_lag = gleam@float:absolute_value(Lag),
    stationary_variance(Params) * math:exp(
        +0.0 - (erlang:element(2, Params) * Abs_lag)
    ).

-file("src/viva_math/ou.gleam", 233).
?DOC(
    " Half-life of mean reversion: `ln(2) / θ`.\n"
    "\n"
    " Time at which `E[X_t]` has covered half the gap toward `μ`.\n"
).
-spec half_life(o_u_params1_d()) -> float().
half_life(Params) ->
    case viva_math@scalar:logarithm(2.0) of
        {ok, L} ->
            case erlang:element(2, Params) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> L / Gleam@denominator
            end;

        {error, _} ->
            +0.0
    end.

-file("src/viva_math/ou.gleam", 246).
?DOC(
    " One Vec3 OU step. Each axis (P, A, D) updated independently via the exact\n"
    " 1D kernel. Three normals drawn from the seed in sequence.\n"
).
-spec step_vec3(
    o_u_params_vec3(),
    viva_math@vector:vec3(),
    float(),
    viva_math@random:seed()
) -> {viva_math@vector:vec3(), viva_math@random:seed()}.
step_vec3(Params, X, Dt, Seed) ->
    {o_u_params_vec3, Th, Mu, Sg} = Params,
    {Px, S1} = step(
        {o_u_params1_d,
            erlang:element(2, Th),
            erlang:element(2, Mu),
            erlang:element(2, Sg)},
        erlang:element(2, X),
        Dt,
        Seed
    ),
    {Py, S2} = step(
        {o_u_params1_d,
            erlang:element(3, Th),
            erlang:element(3, Mu),
            erlang:element(3, Sg)},
        erlang:element(3, X),
        Dt,
        S1
    ),
    {Pz, S3} = step(
        {o_u_params1_d,
            erlang:element(4, Th),
            erlang:element(4, Mu),
            erlang:element(4, Sg)},
        erlang:element(4, X),
        Dt,
        S2
    ),
    {{vec3, Px, Py, Pz}, S3}.

-file("src/viva_math/ou.gleam", 297).
-spec simulate_vec3_loop(
    kernel(),
    kernel(),
    kernel(),
    viva_math@vector:vec3(),
    viva_math@vector:vec3(),
    integer(),
    viva_math@random:seed(),
    list(viva_math@vector:vec3())
) -> {list(viva_math@vector:vec3()), viva_math@random:seed()}.
simulate_vec3_loop(Kx, Ky, Kz, Mu, X, N, Seed, Acc) ->
    case N =< 0 of
        true ->
            {Acc, Seed};

        false ->
            {Zx, S1} = viva_math@random:standard_normal(Seed),
            {Zy, S2} = viva_math@random:standard_normal(S1),
            {Zz, S3} = viva_math@random:standard_normal(S2),
            Nx = (erlang:element(2, Mu) + ((erlang:element(2, X) - erlang:element(
                2,
                Mu
            ))
            * erlang:element(2, Kx)))
            + (erlang:element(3, Kx) * Zx),
            Ny = (erlang:element(3, Mu) + ((erlang:element(3, X) - erlang:element(
                3,
                Mu
            ))
            * erlang:element(2, Ky)))
            + (erlang:element(3, Ky) * Zy),
            Nz = (erlang:element(4, Mu) + ((erlang:element(4, X) - erlang:element(
                4,
                Mu
            ))
            * erlang:element(2, Kz)))
            + (erlang:element(3, Kz) * Zz),
            X_next = {vec3, Nx, Ny, Nz},
            simulate_vec3_loop(
                Kx,
                Ky,
                Kz,
                Mu,
                X_next,
                N - 1,
                S3,
                [X_next | Acc]
            )
    end.

-file("src/viva_math/ou.gleam", 283).
-spec build_kernel(float(), float(), float()) -> kernel().
build_kernel(Theta, Sigma, Dt) ->
    Decay = math:exp(+0.0 - (Theta * Dt)),
    Var_term = case (2.0 * Theta) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> (Sigma * Sigma) * (+0.0 - viva_math@scalar:expm1(
            +0.0 - ((2.0 * Theta) * Dt)
        ))
        / Gleam@denominator
    end,
    Std = case gleam@float:square_root(Var_term) of
        {ok, S} ->
            S;

        {error, _} ->
            +0.0
    end,
    {kernel, Decay, Std}.

-file("src/viva_math/ou.gleam", 263).
?DOC(
    " Simulate Vec3 trajectory. Returns `n` Vec3 points excluding initial.\n"
    "\n"
    " Pre-computes the three componentwise transition kernels once — the loop\n"
    " only does multiply-adds and three normal draws per step.\n"
).
-spec simulate_vec3(
    o_u_params_vec3(),
    viva_math@vector:vec3(),
    float(),
    integer(),
    viva_math@random:seed()
) -> {list(viva_math@vector:vec3()), viva_math@random:seed()}.
simulate_vec3(Params, X0, Dt, N, Seed) ->
    {o_u_params_vec3, Th, Mu, Sg} = Params,
    Kx = build_kernel(erlang:element(2, Th), erlang:element(2, Sg), Dt),
    Ky = build_kernel(erlang:element(3, Th), erlang:element(3, Sg), Dt),
    Kz = build_kernel(erlang:element(4, Th), erlang:element(4, Sg), Dt),
    {Traj, S} = simulate_vec3_loop(Kx, Ky, Kz, Mu, X0, N, Seed, []),
    {lists:reverse(Traj), S}.

-file("src/viva_math/ou.gleam", 323).
?DOC(" Closed-form `E[X_t | X_0 = x0]` componentwise.\n").
-spec mean_at_vec3(o_u_params_vec3(), viva_math@vector:vec3(), float()) -> viva_math@vector:vec3().
mean_at_vec3(Params, X0, T) ->
    {vec3,
        mean_at(
            {o_u_params1_d,
                erlang:element(2, erlang:element(2, Params)),
                erlang:element(2, erlang:element(3, Params)),
                erlang:element(2, erlang:element(4, Params))},
            erlang:element(2, X0),
            T
        ),
        mean_at(
            {o_u_params1_d,
                erlang:element(3, erlang:element(2, Params)),
                erlang:element(3, erlang:element(3, Params)),
                erlang:element(3, erlang:element(4, Params))},
            erlang:element(3, X0),
            T
        ),
        mean_at(
            {o_u_params1_d,
                erlang:element(4, erlang:element(2, Params)),
                erlang:element(4, erlang:element(3, Params)),
                erlang:element(4, erlang:element(4, Params))},
            erlang:element(4, X0),
            T
        )}.

-file("src/viva_math/ou.gleam", 332).
?DOC(" Closed-form `Var[X_t | X_0 = x0]` componentwise.\n").
-spec variance_at_vec3(o_u_params_vec3(), viva_math@vector:vec3(), float()) -> viva_math@vector:vec3().
variance_at_vec3(Params, X0, T) ->
    {vec3,
        variance_at(
            {o_u_params1_d,
                erlang:element(2, erlang:element(2, Params)),
                erlang:element(2, erlang:element(3, Params)),
                erlang:element(2, erlang:element(4, Params))},
            erlang:element(2, X0),
            T
        ),
        variance_at(
            {o_u_params1_d,
                erlang:element(3, erlang:element(2, Params)),
                erlang:element(3, erlang:element(3, Params)),
                erlang:element(3, erlang:element(4, Params))},
            erlang:element(3, X0),
            T
        ),
        variance_at(
            {o_u_params1_d,
                erlang:element(4, erlang:element(2, Params)),
                erlang:element(4, erlang:element(3, Params)),
                erlang:element(4, erlang:element(4, Params))},
            erlang:element(4, X0),
            T
        )}.

-file("src/viva_math/ou.gleam", 353).
?DOC(" Stationary variance per axis.\n").
-spec stationary_variance_vec3(o_u_params_vec3()) -> viva_math@vector:vec3().
stationary_variance_vec3(Params) ->
    {vec3,
        stationary_variance(
            {o_u_params1_d,
                erlang:element(2, erlang:element(2, Params)),
                erlang:element(2, erlang:element(3, Params)),
                erlang:element(2, erlang:element(4, Params))}
        ),
        stationary_variance(
            {o_u_params1_d,
                erlang:element(3, erlang:element(2, Params)),
                erlang:element(3, erlang:element(3, Params)),
                erlang:element(3, erlang:element(4, Params))}
        ),
        stationary_variance(
            {o_u_params1_d,
                erlang:element(4, erlang:element(2, Params)),
                erlang:element(4, erlang:element(3, Params)),
                erlang:element(4, erlang:element(4, Params))}
        )}.