viva_math

-module(viva_math@distributions).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_math/distributions.gleam").
-export([standard_normal/0, gaussian_pdf/2, gaussian_log_pdf/2, gaussian_cdf/2, gaussian_sample/2, uniform_pdf/2, uniform_cdf/2, uniform_sample/2, exponential_pdf/2, exponential_cdf/2, exponential_sample/2, laplace_pdf/2, laplace_sample/2, cauchy_pdf/2, cauchy_sample/2, bernoulli_pmf/2, bernoulli_sample/2, categorical_pmf/2, categorical_entropy/1, categorical_sample/2]).
-export_type([gaussian/0, uniform/0, exponential/0, laplace/0, cauchy/0, bernoulli/0, categorical/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(
    " Probability distributions.\n"
    "\n"
    " Closed-form PDF / CDF / log-PDF + samplers backed by `viva_math/random`.\n"
    " All samplers are pure: they take a `Seed` and return a new seed.\n"
    "\n"
    " ## Scope\n"
    "\n"
    " Continuous: `gaussian`, `uniform`, `exponential`, `cauchy`, `laplace`.\n"
    " Discrete: `bernoulli`, `categorical`.\n"
    "\n"
    " For heavier statistics (gamma, beta, chi-square) prefer `gleam_stats`\n"
    " or `gleam_community_maths`. Here we ship only what `viva_emotion` and\n"
    " `viva_tensor` consume directly.\n"
).

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

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

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

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

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

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

-type categorical() :: {categorical, list(float())}.

-file("src/viva_math/distributions.gleam", 30).
?DOC(" Standard normal N(0, 1).\n").
-spec standard_normal() -> gaussian().
standard_normal() ->
    {gaussian, +0.0, 1.0}.

-file("src/viva_math/distributions.gleam", 35).
?DOC(" Probability density f(x; μ, σ) = 1/(σ√(2π)) · exp(-(x-μ)²/(2σ²)).\n").
-spec gaussian_pdf(gaussian(), float()) -> float().
gaussian_pdf(G, X) ->
    Z = case erlang:element(3, G) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> (X - erlang:element(2, G)) / Gleam@denominator
    end,
    Coeff = case erlang:element(3, G) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator@1 -> 0.3989422804014327 / Gleam@denominator@1
    end,
    Coeff * math:exp((-0.5 * Z) * Z).

-file("src/viva_math/distributions.gleam", 42).
?DOC(" Log-density. More numerically stable than `ln(gaussian_pdf)`.\n").
-spec gaussian_log_pdf(gaussian(), float()) -> float().
gaussian_log_pdf(G, X) ->
    Z = case erlang:element(3, G) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> (X - erlang:element(2, G)) / Gleam@denominator
    end,
    (((-0.5 * Z) * Z) - math:log(erlang:element(3, G))) - (0.5 * math:log(
        6.283185307179586
    )).

-file("src/viva_math/distributions.gleam", 48).
?DOC(" Cumulative distribution F(x; μ, σ) = ½ · (1 + erf((x - μ)/(σ√2))).\n").
-spec gaussian_cdf(gaussian(), float()) -> float().
gaussian_cdf(G, X) ->
    0.5 * (1.0 + math:erf(case (erlang:element(3, G) * 1.4142135623730951) of
            +0.0 -> +0.0;
            -0.0 -> -0.0;
            Gleam@denominator -> (X - erlang:element(2, G)) / Gleam@denominator
        end)).

-file("src/viva_math/distributions.gleam", 54).
?DOC(" Sample from a Gaussian.\n").
-spec gaussian_sample(gaussian(), viva_math@random:seed()) -> {float(),
    viva_math@random:seed()}.
gaussian_sample(G, Seed) ->
    viva_math@random:normal(Seed, erlang:element(2, G), erlang:element(3, G)).

-file("src/viva_math/distributions.gleam", 70).
?DOC(" PDF on [low, high]; zero elsewhere.\n").
-spec uniform_pdf(uniform(), float()) -> float().
uniform_pdf(U, X) ->
    case (X < erlang:element(2, U)) orelse (X > erlang:element(3, U)) of
        true ->
            +0.0;

        false ->
            case (erlang:element(3, U) - erlang:element(2, U)) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> 1.0 / Gleam@denominator
            end
    end.

-file("src/viva_math/distributions.gleam", 78).
?DOC(" CDF.\n").
-spec uniform_cdf(uniform(), float()) -> float().
uniform_cdf(U, X) ->
    case {X < erlang:element(2, U), X > erlang:element(3, U)} of
        {true, _} ->
            +0.0;

        {_, true} ->
            1.0;

        {_, _} ->
            case (erlang:element(3, U) - erlang:element(2, U)) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> (X - erlang:element(2, U)) / Gleam@denominator
            end
    end.

-file("src/viva_math/distributions.gleam", 86).
-spec uniform_sample(uniform(), viva_math@random:seed()) -> {float(),
    viva_math@random:seed()}.
uniform_sample(U, Seed) ->
    viva_math@random:uniform_in(
        Seed,
        erlang:element(2, U),
        erlang:element(3, U)
    ).

-file("src/viva_math/distributions.gleam", 99).
?DOC(" PDF f(x; λ) = λ · e^(-λx) for x ≥ 0.\n").
-spec exponential_pdf(exponential(), float()) -> float().
exponential_pdf(E, X) ->
    case X < +0.0 of
        true ->
            +0.0;

        false ->
            erlang:element(2, E) * math:exp(+0.0 - (erlang:element(2, E) * X))
    end.

-file("src/viva_math/distributions.gleam", 107).
?DOC(" CDF.\n").
-spec exponential_cdf(exponential(), float()) -> float().
exponential_cdf(E, X) ->
    case X < +0.0 of
        true ->
            +0.0;

        false ->
            1.0 - math:exp(+0.0 - (erlang:element(2, E) * X))
    end.

-file("src/viva_math/distributions.gleam", 119).
?DOC(
    " Inverse-CDF sampling.\n"
    "\n"
    " Uses x = -log1p(-U) / λ, which is numerically equivalent to\n"
    " -ln(1 - U) / λ but avoids cancellation when U is near zero. Since\n"
    " `random.uniform` returns U ∈ [0, 1), 1 - U ∈ (0, 1] never hits zero.\n"
).
-spec exponential_sample(exponential(), viva_math@random:seed()) -> {float(),
    viva_math@random:seed()}.
exponential_sample(E, Seed) ->
    {U, S} = viva_math@random:uniform(Seed),
    {+0.0 - (case erlang:element(2, E) of
            +0.0 -> +0.0;
            -0.0 -> -0.0;
            Gleam@denominator -> viva_math@scalar:log1p(+0.0 - U) / Gleam@denominator
        end), S}.

-file("src/viva_math/distributions.gleam", 135).
-spec laplace_pdf(laplace(), float()) -> float().
laplace_pdf(L, X) ->
    Abs_dev = case X >= erlang:element(2, L) of
        true ->
            X - erlang:element(2, L);

        false ->
            erlang:element(2, L) - X
    end,
    case (2.0 * erlang:element(3, L)) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator@1 -> math:exp(+0.0 - (case erlang:element(3, L) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> Abs_dev / Gleam@denominator
            end)) / Gleam@denominator@1
    end.

-file("src/viva_math/distributions.gleam", 143).
-spec laplace_sample(laplace(), viva_math@random:seed()) -> {float(),
    viva_math@random:seed()}.
laplace_sample(L, Seed) ->
    {U, S} = viva_math@random:uniform(Seed),
    Shifted = U - 0.5,
    Abs_shift = case Shifted >= +0.0 of
        true ->
            Shifted;

        false ->
            +0.0 - Shifted
    end,
    Sign_shift = viva_math@scalar:sign(Shifted),
    Inner = +0.0 - (2.0 * Abs_shift),
    Log_factor = case Inner =< -1.0 of
        true ->
            +0.0 - 1.0e300;

        false ->
            viva_math@scalar:log1p(Inner)
    end,
    {erlang:element(2, L) - ((erlang:element(3, L) * Sign_shift) * Log_factor),
        S}.

-file("src/viva_math/distributions.gleam", 170).
-spec cauchy_pdf(cauchy(), float()) -> float().
cauchy_pdf(C, X) ->
    Z = case erlang:element(3, C) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> (X - erlang:element(2, C)) / Gleam@denominator
    end,
    case ((3.141592653589793 * erlang:element(3, C)) * (1.0 + (Z * Z))) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator@1 -> 1.0 / Gleam@denominator@1
    end.

-file("src/viva_math/distributions.gleam", 175).
-spec cauchy_sample(cauchy(), viva_math@random:seed()) -> {float(),
    viva_math@random:seed()}.
cauchy_sample(C, Seed) ->
    {U, S} = viva_math@random:uniform(Seed),
    Arg = 3.141592653589793 * (U - 0.5),
    {erlang:element(2, C) + (erlang:element(3, C) * math:tan(Arg)), S}.

-file("src/viva_math/distributions.gleam", 190).
-spec bernoulli_pmf(bernoulli(), integer()) -> float().
bernoulli_pmf(B, K) ->
    case K of
        1 ->
            erlang:element(2, B);

        0 ->
            1.0 - erlang:element(2, B);

        _ ->
            +0.0
    end.

-file("src/viva_math/distributions.gleam", 198).
-spec bernoulli_sample(bernoulli(), viva_math@random:seed()) -> {boolean(),
    viva_math@random:seed()}.
bernoulli_sample(B, Seed) ->
    viva_math@random:bernoulli(Seed, erlang:element(2, B)).

-file("src/viva_math/distributions.gleam", 242).
-spec list_at(list(float()), integer()) -> {ok, float()} | {error, nil}.
list_at(Xs, Idx) ->
    case {Xs, Idx} of
        {[], _} ->
            {error, nil};

        {[X | _], 0} ->
            {ok, X};

        {[_ | Rest], N} ->
            list_at(Rest, N - 1)
    end.

-file("src/viva_math/distributions.gleam", 214).
?DOC(" PMF for a category index (zero-based).\n").
-spec categorical_pmf(categorical(), integer()) -> float().
categorical_pmf(C, K) ->
    case list_at(erlang:element(2, C), K) of
        {ok, P} ->
            P;

        {error, _} ->
            +0.0
    end.

-file("src/viva_math/distributions.gleam", 222).
?DOC(" Entropy H(X) = -Σ pᵢ log pᵢ.\n").
-spec categorical_entropy(categorical()) -> float().
categorical_entropy(C) ->
    gleam@list:fold(erlang:element(2, C), +0.0, fun(Acc, P) -> case P =< +0.0 of
                true ->
                    Acc;

                false ->
                    Acc - (P * math:log(P))
            end end).

-file("src/viva_math/distributions.gleam", 231).
-spec categorical_sample(categorical(), viva_math@random:seed()) -> {ok,
        {integer(), viva_math@random:seed()}} |
    {error, nil}.
categorical_sample(C, Seed) ->
    viva_math@random:categorical(Seed, erlang:element(2, C)).