viva_math

-module(viva_math@precision).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_math/precision.gleam").
-export([neumaier_sum/1, kahan_sum/1, pairwise_sum/1, two_sum/2, fsum/1, neumaier_mean/1, moments_empty/0, moments_update/2, moments_from_list/1, moments_mean/1, moments_variance/1, moments_sample_variance/1, moments_skewness/1, moments_excess_kurtosis/1, moments_combine/2, cmp_abs_desc/2]).
-export_type([moments/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(
    " High-precision numerical primitives.\n"
    "\n"
    " IEEE-754 doubles lose precision when summing values that differ widely\n"
    " in magnitude. This module provides compensated and exact summation\n"
    " algorithms plus Pébay's online higher-order moment accumulators.\n"
    "\n"
    " ## Algorithms shipped\n"
    "\n"
    " | Function           | Worst-case error | Cost   | When to use                |\n"
    " | ------------------ | ---------------- | ------ | -------------------------- |\n"
    " | `neumaier_sum`     | O(ε)             | ~3x    | **Default — fast + safe**  |\n"
    " | `kahan_sum`        | O(ε)             | ~4x    | Reference / didactic       |\n"
    " | `pairwise_sum`     | O(log n · ε)     | ~1x    | Reductions on long lists   |\n"
    " | `fsum`             | round-once exact | ~10x   | Critical accuracy           |\n"
    "\n"
    " `neumaier_sum` is what CPython 3.12 `sum()` uses by default and matches\n"
    " the Gonum and Julia recommendations. Catastrophic example:\n"
    "\n"
    " ```gleam\n"
    " import viva_math/precision\n"
    "\n"
    " precision.neumaier_sum([1.0, 1.0e100, 1.0, -1.0e100])\n"
    " // -> 2.0  (naive sum would give 0.0)\n"
    " ```\n"
    "\n"
    " ## References\n"
    "\n"
    " - Kahan (1965) \"Pracniques: further remarks on reducing truncation errors\"\n"
    " - Neumaier (1974) \"Rundungsfehleranalyse einiger Verfahren zur Summation\n"
    "   endlicher Summen\"\n"
    " - Shewchuk (1997) \"Adaptive Precision Floating-Point Arithmetic\"\n"
    " - Pébay (2008) \"Formulas for robust, one-pass parallel computation of\n"
    "   covariances and arbitrary-order statistical moments\" (Sandia)\n"
    " - Pébay, Terriberry et al. (2016) \"Numerically Stable, Scalable Formulas\n"
    "   for Parallel and Online Computation of Higher-Order Multivariate\n"
    "   Central Moments with Arbitrary Weights\"\n"
    " - CPython issue #100425 (Python 3.12 switched to Neumaier)\n"
).

-type moments() :: {moments, integer(), float(), float(), float(), float()}.

-file("src/viva_math/precision.gleam", 63).
-spec neumaier_loop(list(float()), float(), float()) -> float().
neumaier_loop(Xs, Sum, Comp) ->
    case Xs of
        [] ->
            Sum + Comp;

        [X | Rest] ->
            T = Sum + X,
            New_comp = case gleam@float:absolute_value(Sum) >= gleam@float:absolute_value(
                X
            ) of
                true ->
                    Comp + ((Sum - T) + X);

                false ->
                    Comp + ((X - T) + Sum)
            end,
            neumaier_loop(Rest, T, New_comp)
    end.

-file("src/viva_math/precision.gleam", 59).
?DOC(
    " Neumaier (Kahan-Babuška improved) compensated sum.\n"
    "\n"
    " Maintains a separate compensation accumulator and swaps the role of the\n"
    " running sum / next term depending on magnitude. Recovers about 16 extra\n"
    " bits of precision over naive summation at ~3x the cost.\n"
    "\n"
    " ## Example\n"
    "\n"
    " ```gleam\n"
    " precision.neumaier_sum([1.0, 1.0e100, 1.0, -1.0e100])\n"
    " // -> 2.0   (naive sum returns 0.0)\n"
    " ```\n"
).
-spec neumaier_sum(list(float())) -> float().
neumaier_sum(Xs) ->
    neumaier_loop(Xs, +0.0, +0.0).

-file("src/viva_math/precision.gleam", 88).
-spec kahan_loop(list(float()), float(), float()) -> float().
kahan_loop(Xs, Sum, Comp) ->
    case Xs of
        [] ->
            Sum;

        [X | Rest] ->
            Y = X - Comp,
            T = Sum + Y,
            New_comp = (T - Sum) - Y,
            kahan_loop(Rest, T, New_comp)
    end.

-file("src/viva_math/precision.gleam", 84).
?DOC(
    " Classical Kahan compensated sum.\n"
    "\n"
    " Slightly less accurate than `neumaier_sum` (fails on the pathological\n"
    " `[1, 1e100, 1, -1e100]` example) but easier to reason about. Useful as a\n"
    " reference implementation.\n"
).
-spec kahan_sum(list(float())) -> float().
kahan_sum(Xs) ->
    kahan_loop(Xs, +0.0, +0.0).

-file("src/viva_math/precision.gleam", 105).
?DOC(
    " Pairwise sum: recursively splits and combines.\n"
    "\n"
    " Error grows as O(log n · ε) instead of O(n · ε) for naive sum, at the\n"
    " same asymptotic cost. NumPy's default. Good for very long lists where\n"
    " the constant factor of Neumaier matters.\n"
).
-spec pairwise_sum(list(float())) -> float().
pairwise_sum(Xs) ->
    case Xs of
        [] ->
            +0.0;

        [X] ->
            X;

        [A, B] ->
            A + B;

        _ ->
            N = erlang:length(Xs),
            Mid = N div 2,
            Left = gleam@list:take(Xs, Mid),
            Right = gleam@list:drop(Xs, Mid),
            pairwise_sum(Left) + pairwise_sum(Right)
    end.

-file("src/viva_math/precision.gleam", 157).
?DOC(
    " Two-sum: returns the exact sum `a+b` as a non-overlapping pair `(hi, lo)`\n"
    " where `hi = round(a+b)` and `lo = (a+b) - hi`.\n"
).
-spec two_sum(float(), float()) -> {float(), float()}.
two_sum(A, B) ->
    S = A + B,
    Bp = S - A,
    Ap = S - Bp,
    Lo = (A - Ap) + (B - Bp),
    {S, Lo}.

-file("src/viva_math/precision.gleam", 134).
-spec fsum_merge(list(float()), float(), list(float())) -> list(float()).
fsum_merge(Partials, X, Acc) ->
    case Partials of
        [] ->
            case X =:= +0.0 of
                true ->
                    lists:reverse(Acc);

                false ->
                    lists:reverse([X | Acc])
            end;

        [P | Rest] ->
            {Hi, Lo} = two_sum(P, X),
            case Lo =:= +0.0 of
                true ->
                    fsum_merge(Rest, Hi, Acc);

                false ->
                    fsum_merge(Rest, Hi, [Lo | Acc])
            end
    end.

-file("src/viva_math/precision.gleam", 130).
-spec fsum_add(list(float()), float()) -> list(float()).
fsum_add(Partials, X) ->
    fsum_merge(Partials, X, []).

-file("src/viva_math/precision.gleam", 125).
?DOC(
    " Shewchuk's fsum: maintains a list of non-overlapping partial sums.\n"
    "\n"
    " Round-to-nearest exact result up to a final rounding. Matches Python's\n"
    " `math.fsum`. The most accurate option available, at ~10x the cost of\n"
    " naive summation.\n"
).
-spec fsum(list(float())) -> float().
fsum(Xs) ->
    Partials = gleam@list:fold(Xs, [], fun fsum_add/2),
    gleam@list:fold(Partials, +0.0, fun(Acc, P) -> Acc + P end).

-file("src/viva_math/precision.gleam", 170).
?DOC(" Mean using Neumaier-compensated sum. Errors on empty input.\n").
-spec neumaier_mean(list(float())) -> {ok, float()} | {error, nil}.
neumaier_mean(Xs) ->
    case Xs of
        [] ->
            {error, nil};

        _ ->
            {ok, case erlang:float(erlang:length(Xs)) of
                    +0.0 -> +0.0;
                    -0.0 -> -0.0;
                    Gleam@denominator -> neumaier_sum(Xs) / Gleam@denominator
                end}
    end.

-file("src/viva_math/precision.gleam", 191).
?DOC(" Empty accumulator.\n").
-spec moments_empty() -> moments().
moments_empty() ->
    {moments, 0, +0.0, +0.0, +0.0, +0.0}.

-file("src/viva_math/precision.gleam", 204).
?DOC(
    " Update with a single sample (Pébay 2008 recurrence).\n"
    "\n"
    " Let n be the new count, δ = x - mean, δ_n = δ/n, δ_n² and term1 = δ·δ_n·(n-1).\n"
    " Then:\n"
    "   M₂ ← M₂ + term1\n"
    "   M₃ ← M₃ + term1·δ_n·(n-2) - 3·δ_n·M₂\n"
    "   M₄ ← M₄ + term1·δ_n²·(n²-3n+3) + 6·δ_n²·M₂ - 4·δ_n·M₃\n"
    "\n"
    " Updates the mean last to preserve the previous-iteration moments above.\n"
).
-spec moments_update(moments(), float()) -> moments().
moments_update(M, X) ->
    N1 = erlang:element(2, M) + 1,
    N1_float = erlang:float(N1),
    Delta = X - erlang:element(3, M),
    Delta_n = case N1_float of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> Delta / Gleam@denominator
    end,
    Delta_n_sq = Delta_n * Delta_n,
    N_prev = erlang:element(2, M),
    N_prev_float = erlang:float(N_prev),
    Term1 = (Delta * Delta_n) * N_prev_float,
    New_m4 = ((erlang:element(6, M) + ((Term1 * Delta_n_sq) * (((N1_float * N1_float)
    - (3.0 * N1_float))
    + 3.0)))
    + ((6.0 * Delta_n_sq) * erlang:element(4, M)))
    - ((4.0 * Delta_n) * erlang:element(5, M)),
    New_m3 = (erlang:element(5, M) + ((Term1 * Delta_n) * (N1_float - 2.0))) - ((3.0
    * Delta_n)
    * erlang:element(4, M)),
    New_m2 = erlang:element(4, M) + Term1,
    New_mean = erlang:element(3, M) + Delta_n,
    {moments, N1, New_mean, New_m2, New_m3, New_m4}.

-file("src/viva_math/precision.gleam", 236).
?DOC(" Build accumulator from a list.\n").
-spec moments_from_list(list(float())) -> moments().
moments_from_list(Xs) ->
    gleam@list:fold(Xs, moments_empty(), fun moments_update/2).

-file("src/viva_math/precision.gleam", 241).
?DOC(" Population mean.\n").
-spec moments_mean(moments()) -> {ok, float()} | {error, nil}.
moments_mean(M) ->
    case erlang:element(2, M) of
        0 ->
            {error, nil};

        _ ->
            {ok, erlang:element(3, M)}
    end.

-file("src/viva_math/precision.gleam", 249).
?DOC(" Population variance σ² = M₂ / n.\n").
-spec moments_variance(moments()) -> {ok, float()} | {error, nil}.
moments_variance(M) ->
    case erlang:element(2, M) of
        0 ->
            {error, nil};

        _ ->
            {ok, case erlang:float(erlang:element(2, M)) of
                    +0.0 -> +0.0;
                    -0.0 -> -0.0;
                    Gleam@denominator -> erlang:element(4, M) / Gleam@denominator
                end}
    end.

-file("src/viva_math/precision.gleam", 257).
?DOC(" Sample variance s² = M₂ / (n-1).\n").
-spec moments_sample_variance(moments()) -> {ok, float()} | {error, nil}.
moments_sample_variance(M) ->
    case erlang:element(2, M) of
        0 ->
            {error, nil};

        1 ->
            {error, nil};

        _ ->
            {ok, case erlang:float(erlang:element(2, M) - 1) of
                    +0.0 -> +0.0;
                    -0.0 -> -0.0;
                    Gleam@denominator -> erlang:element(4, M) / Gleam@denominator
                end}
    end.

-file("src/viva_math/precision.gleam", 265).
?DOC(" Population skewness γ₁ = (M₃/n) / (M₂/n)^(3/2) = √n · M₃ / M₂^(3/2).\n").
-spec moments_skewness(moments()) -> {ok, float()} | {error, nil}.
moments_skewness(M) ->
    case erlang:element(2, M) of
        0 ->
            {error, nil};

        N ->
            case erlang:element(4, M) =< +0.0 of
                true ->
                    {error, nil};

                false ->
                    N_float = erlang:float(N),
                    M2_sqrt = math:sqrt(erlang:element(4, M)),
                    M2_pow_3_2 = M2_sqrt * erlang:element(4, M),
                    {ok, case M2_pow_3_2 of
                            +0.0 -> +0.0;
                            -0.0 -> -0.0;
                            Gleam@denominator -> math:sqrt(N_float) * erlang:element(
                                5,
                                M
                            )
                            / Gleam@denominator
                        end}
            end
    end.

-file("src/viva_math/precision.gleam", 284).
?DOC(" Excess kurtosis γ₂ = n · M₄ / M₂² - 3.\n").
-spec moments_excess_kurtosis(moments()) -> {ok, float()} | {error, nil}.
moments_excess_kurtosis(M) ->
    case erlang:element(2, M) of
        0 ->
            {error, nil};

        N ->
            case erlang:element(4, M) =< +0.0 of
                true ->
                    {error, nil};

                false ->
                    N_float = erlang:float(N),
                    {ok, (case (erlang:element(4, M) * erlang:element(4, M)) of
                            +0.0 -> +0.0;
                            -0.0 -> -0.0;
                            Gleam@denominator -> N_float * erlang:element(6, M)
                            / Gleam@denominator
                        end) - 3.0}
            end
    end.

-file("src/viva_math/precision.gleam", 303).
?DOC(
    " Combine two Pébay accumulators computed in parallel (Chan formula).\n"
    "\n"
    " Useful for distributed / parallel reductions: split a stream across\n"
    " workers, then merge.\n"
).
-spec moments_combine(moments(), moments()) -> moments().
moments_combine(A, B) ->
    case {erlang:element(2, A), erlang:element(2, B)} of
        {0, _} ->
            B;

        {_, 0} ->
            A;

        {Na, Nb} ->
            Na_f = erlang:float(Na),
            Nb_f = erlang:float(Nb),
            N = Na + Nb,
            N_f = erlang:float(N),
            Delta = erlang:element(3, B) - erlang:element(3, A),
            Delta2 = Delta * Delta,
            Delta3 = Delta2 * Delta,
            Delta4 = Delta2 * Delta2,
            New_mean = erlang:element(3, A) + (case N_f of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> Delta * Nb_f / Gleam@denominator
            end),
            New_m2 = (erlang:element(4, A) + erlang:element(4, B)) + (case N_f of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator@1 -> (Delta2 * Na_f) * Nb_f / Gleam@denominator@1
            end),
            New_m3 = ((erlang:element(5, A) + erlang:element(5, B)) + (case (N_f
            * N_f) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator@2 -> ((Delta3 * Na_f) * Nb_f) * (Na_f - Nb_f)
                / Gleam@denominator@2
            end))
            + (case N_f of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator@3 -> (3.0 * Delta) * ((Na_f * erlang:element(
                    4,
                    B
                ))
                - (Nb_f * erlang:element(4, A)))
                / Gleam@denominator@3
            end),
            New_m4 = (((erlang:element(6, A) + erlang:element(6, B)) + (case ((N_f
            * N_f)
            * N_f) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator@4 -> ((Delta4 * Na_f) * Nb_f) * (((Na_f * Na_f)
                - (Na_f * Nb_f))
                + (Nb_f * Nb_f))
                / Gleam@denominator@4
            end))
            + (case (N_f * N_f) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator@5 -> (6.0 * Delta2) * (((Na_f * Na_f) * erlang:element(
                    4,
                    B
                ))
                + ((Nb_f * Nb_f) * erlang:element(4, A)))
                / Gleam@denominator@5
            end))
            + (case N_f of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator@6 -> (4.0 * Delta) * ((Na_f * erlang:element(
                    5,
                    B
                ))
                - (Nb_f * erlang:element(5, A)))
                / Gleam@denominator@6
            end),
            {moments, N, New_mean, New_m2, New_m3, New_m4}
    end.

-file("src/viva_math/precision.gleam", 365).
?DOC(
    " Stable cmp by absolute value (descending). Useful for sorting before\n"
    " summation in worst-case scenarios.\n"
).
-spec cmp_abs_desc(float(), float()) -> gleam@order:order().
cmp_abs_desc(A, B) ->
    case {gleam@float:absolute_value(A) > gleam@float:absolute_value(B),
        gleam@float:absolute_value(A) < gleam@float:absolute_value(B)} of
        {true, _} ->
            lt;

        {_, true} ->
            gt;

        {_, _} ->
            eq
    end.