viva_math

-module(viva_math@special).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_math/special.gleam").
-export([lgamma/1, gamma/1, beta/2, lbeta/2, digamma/1, factorial/1, binomial/2]).

-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(
    " Special mathematical functions.\n"
    "\n"
    " Implementations focus on functions that complement `gleam_community_maths`\n"
    " and are needed by `viva_math/distributions` for Gamma/Beta/Dirichlet\n"
    " distributions, by `viva_math/free_energy` for KL divergences of\n"
    " exponential-family distributions, and by `viva_math/entropy` for the\n"
    " differential entropy of those families.\n"
    "\n"
    " ## Algorithms\n"
    "\n"
    " - `gamma`  → Lanczos approximation (g=7, 9 coefficients). Accurate to\n"
    "              ~15 decimal digits across the entire positive reals.\n"
    " - `lgamma` → Logarithmic form of Lanczos; preferred numerically because\n"
    "              the regular gamma overflows past Γ(171.6) ≈ 10³⁰⁸.\n"
    " - `digamma`→ ψ(x) via the asymptotic series for large x with the\n"
    "              recurrence ψ(x) = ψ(x+1) - 1/x for small x.\n"
    " - `beta`   → exp(lgamma(x) + lgamma(y) - lgamma(x+y)).\n"
    " - `factorial` → Stirling-style via gamma for large n, exact for n ≤ 20.\n"
    "\n"
    " ## References\n"
    "\n"
    " - Lanczos (1964) \"A precision approximation of the gamma function\"\n"
    " - Press, Teukolsky et al. (1992) \"Numerical Recipes\" §6.1\n"
    " - GSL `gsl_sf_lngamma`\n"
).

-file("src/viva_math/special.gleam", 66).
-spec lanczos_series(float()) -> float().
lanczos_series(Xm1) ->
    (((((((0.99999999999980993 + (case (Xm1 + 1.0) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> 676.5203681218851 / Gleam@denominator
    end)) - (case (Xm1 + 2.0) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator@1 -> 1259.1392167224028 / Gleam@denominator@1
    end)) + (case (Xm1 + 3.0) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator@2 -> 771.32342877765313 / Gleam@denominator@2
    end)) - (case (Xm1 + 4.0) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator@3 -> 176.61502916214059 / Gleam@denominator@3
    end)) + (case (Xm1 + 5.0) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator@4 -> 12.507343278686905 / Gleam@denominator@4
    end)) - (case (Xm1 + 6.0) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator@5 -> 0.13857109526572012 / Gleam@denominator@5
    end)) + (case (Xm1 + 7.0) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator@6 -> 9.9843695780195716e-6 / Gleam@denominator@6
    end)) + (case (Xm1 + 8.0) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator@7 -> 1.5056327351493116e-7 / Gleam@denominator@7
    end).

-file("src/viva_math/special.gleam", 58).
-spec lanczos_lgamma(float()) -> float().
lanczos_lgamma(X) ->
    Xm1 = X - 1.0,
    T = Xm1 + 7.5,
    Series = lanczos_series(Xm1),
    (math:log(2.5066282746310002 * Series) + ((Xm1 + 0.5) * math:log(T))) - T.

-file("src/viva_math/special.gleam", 36).
?DOC(
    " Logarithm of the gamma function `ln Γ(x)`. Defined for `x > 0`.\n"
    "\n"
    " Uses the Lanczos approximation with g=7, accurate to ~15 digits.\n"
).
-spec lgamma(float()) -> float().
lgamma(X) ->
    case X =< +0.0 of
        true ->
            1.7976931348623157e308;

        false ->
            case X < 0.5 of
                true ->
                    Sin_pi_x = math:sin(3.141592653589793 * X),
                    Abs_sin = case Sin_pi_x < +0.0 of
                        true ->
                            +0.0 - Sin_pi_x;

                        false ->
                            Sin_pi_x
                    end,
                    (math:log(3.141592653589793) - math:log(Abs_sin)) - lgamma(
                        1.0 - X
                    );

                false ->
                    lanczos_lgamma(X)
            end
    end.

-file("src/viva_math/special.gleam", 88).
?DOC(" Gamma function Γ(x). Overflows past x ≈ 171.6 — prefer `lgamma`.\n").
-spec gamma(float()) -> float().
gamma(X) ->
    math:exp(lgamma(X)).

-file("src/viva_math/special.gleam", 94).
?DOC(
    " Beta function B(x, y) = Γ(x)·Γ(y)/Γ(x+y). Always computed via `lgamma`\n"
    " to avoid overflow.\n"
).
-spec beta(float(), float()) -> float().
beta(X, Y) ->
    math:exp((lgamma(X) + lgamma(Y)) - lgamma(X + Y)).

-file("src/viva_math/special.gleam", 99).
?DOC(" Logarithm of the beta function — `ln B(x, y)`.\n").
-spec lbeta(float(), float()) -> float().
lbeta(X, Y) ->
    (lgamma(X) + lgamma(Y)) - lgamma(X + Y).

-file("src/viva_math/special.gleam", 125).
-spec digamma_asymptotic(float()) -> float().
digamma_asymptotic(X) ->
    Inv_x = case X of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> 1.0 / Gleam@denominator
    end,
    Inv_x2 = Inv_x * Inv_x,
    (math:log(X) - (0.5 * Inv_x)) - (Inv_x2 * (0.0833333333333333333 - (Inv_x2 * (0.00833333333333333333
    - (Inv_x2 * (0.00396825396825396825 - (Inv_x2 * 0.00416666666666666666))))))).

-file("src/viva_math/special.gleam", 118).
-spec digamma_shifted(float(), float()) -> float().
digamma_shifted(X, Accum) ->
    case X < 20.0 of
        true ->
            digamma_shifted(X + 1.0, Accum - (case X of
                    +0.0 -> +0.0;
                    -0.0 -> -0.0;
                    Gleam@denominator -> 1.0 / Gleam@denominator
                end));

        false ->
            Accum + digamma_asymptotic(X)
    end.

-file("src/viva_math/special.gleam", 111).
?DOC(
    " Digamma function ψ(x). Defined for `x > 0`.\n"
    "\n"
    " Combines the recurrence ψ(x) = ψ(x+1) - 1/x to push x ≥ 20, then evaluates\n"
    " the asymptotic series ψ(x) ≈ ln x - 1/(2x) - Σ B₂ₖ / (2k·x^(2k)).\n"
).
-spec digamma(float()) -> float().
digamma(X) ->
    case X =< +0.0 of
        true ->
            +0.0 - 1.7976931348623157e308;

        false ->
            digamma_shifted(X, +0.0)
    end.

-file("src/viva_math/special.gleam", 163).
-spec factorial_exact(integer(), integer()) -> integer().
factorial_exact(N, Acc) ->
    case N =< 1 of
        true ->
            Acc;

        false ->
            factorial_exact(N - 1, Acc * N)
    end.

-file("src/viva_math/special.gleam", 152).
?DOC(
    " Factorial n! for non-negative integer n.\n"
    "\n"
    " Exact for n ≤ 20 (fits in i64). For larger n falls back to Γ(n+1) via\n"
    " Lanczos, with attendant ~15-digit relative accuracy.\n"
).
-spec factorial(integer()) -> {ok, float()} | {error, nil}.
factorial(N) ->
    case N < 0 of
        true ->
            {error, nil};

        false ->
            case N =< 20 of
                true ->
                    {ok, erlang:float(factorial_exact(N, 1))};

                false ->
                    {ok, gamma(erlang:float(N) + 1.0)}
            end
    end.

-file("src/viva_math/special.gleam", 171).
?DOC(" Binomial coefficient C(n, k) = n! / (k!·(n-k)!) via lgamma for stability.\n").
-spec binomial(integer(), integer()) -> {ok, float()} | {error, nil}.
binomial(N, K) ->
    case {N < 0, K < 0, K > N} of
        {true, _, _} ->
            {error, nil};

        {_, true, _} ->
            {error, nil};

        {_, _, true} ->
            {ok, +0.0};

        {_, _, _} ->
            N_f = erlang:float(N),
            K_f = erlang:float(K),
            {ok,
                math:exp(
                    (lgamma(N_f + 1.0) - lgamma(K_f + 1.0)) - lgamma(
                        (N_f - K_f) + 1.0
                    )
                )}
    end.