viva_math

-module(viva_math@tdigest).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_math/tdigest.gleam").
-export([compression/1, centroid_mean/1, centroid_weight/1, with_compression/1, new/0, insert_weighted/3, insert/2, insert_all/2, merge/2, quantile/2, count/1, min/1, max/1, median/1, p99/1]).
-export_type([centroid/0, t_digest/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(
    " t-digest — streaming approximate quantiles.\n"
    "\n"
    " Dunning's t-digest (2014, 2019) provides accurate approximate quantiles\n"
    " over a stream with bounded memory and fast updates. Particularly\n"
    " accurate at the tail (`q < 0.01` or `q > 0.99`), which is where most\n"
    " applications need precision.\n"
    "\n"
    " Memory: O(δ) where δ is the compression parameter (default 100).\n"
    " Update: O(log δ) amortised.\n"
    " Quantile query: O(δ).\n"
    "\n"
    " ## When to use\n"
    "\n"
    " - You need percentiles over a stream that doesn't fit in memory.\n"
    " - Tail quantiles (p99, p99.9) matter — t-digest is exact-ish there.\n"
    " - You want to merge digests from parallel workers.\n"
    "\n"
    " ## Reference\n"
    "\n"
    " Dunning & Ertl (2019) \"Computing Extremely Accurate Quantiles Using\n"
    " t-Digests\" — https://github.com/tdunning/t-digest\n"
    "\n"
    " ## Algorithm summary\n"
    "\n"
    " A t-digest is a sorted set of *centroids*, each (`mean`, `weight`).\n"
    " New samples merge into the nearest centroid up to a size limit dictated\n"
    " by the scale function `k(q) = δ · q · (1 - q) / 2π` — small at the\n"
    " tails, large at the median, so tail centroids stay small and accurate.\n"
).

-opaque centroid() :: {centroid, float(), float()}.

-opaque t_digest() :: {t_digest, float(), list(centroid()), float()}.

-file("src/viva_math/tdigest.gleam", 59).
?DOC(" Compression parameter δ (typical 100; larger → more accurate / more memory).\n").
-spec compression(t_digest()) -> float().
compression(Td) ->
    erlang:element(2, Td).

-file("src/viva_math/tdigest.gleam", 64).
?DOC(" Mean of a centroid.\n").
-spec centroid_mean(centroid()) -> float().
centroid_mean(C) ->
    erlang:element(2, C).

-file("src/viva_math/tdigest.gleam", 69).
?DOC(" Weight of a centroid.\n").
-spec centroid_weight(centroid()) -> float().
centroid_weight(C) ->
    erlang:element(3, C).

-file("src/viva_math/tdigest.gleam", 83).
?DOC(" Empty t-digest with explicit compression parameter.\n").
-spec with_compression(float()) -> t_digest().
with_compression(Delta) ->
    {t_digest, Delta, [], +0.0}.

-file("src/viva_math/tdigest.gleam", 78).
?DOC(" Empty t-digest with default compression (δ = 100).\n").
-spec new() -> t_digest().
new() ->
    with_compression(100.0).

-file("src/viva_math/tdigest.gleam", 297).
?DOC(
    " k(q) scale function: maximum allowed weight for a centroid at cumulative\n"
    " quantile q. Implementation uses the standard scaling\n"
    " `k₁(q) = (δ / 2π) · arcsin(2q - 1)` and translates back to a weight\n"
    " budget. We use the simpler approximation `4 · total · q · (1 - q) / δ`\n"
    " which captures the same \"small at tails\" behaviour with cheaper\n"
    " arithmetic.\n"
).
-spec max_weight_for_quantile(float(), float(), float()) -> float().
max_weight_for_quantile(Q, Delta, Total) ->
    Q_clamped = viva_math@scalar:clamp(Q, +0.0, 1.0),
    case Delta of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> ((4.0 * Total) * Q_clamped) * (1.0 - Q_clamped) / Gleam@denominator
    end.

-file("src/viva_math/tdigest.gleam", 255).
-spec compress_walk(
    list(centroid()),
    centroid(),
    float(),
    float(),
    float(),
    list(centroid())
) -> {list(centroid()), float()}.
compress_walk(Remaining, Current, Current_q_start, Delta, Total, Acc) ->
    case Remaining of
        [] ->
            {[Current | Acc], Current_q_start};

        [Next | Rest] ->
            Combined_weight = erlang:element(3, Current) + erlang:element(
                3,
                Next
            ),
            Q_end = case Total of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> ((Current_q_start * Total) + Combined_weight)
                / Gleam@denominator
            end,
            Allowed = max_weight_for_quantile(Q_end, Delta, Total),
            case Combined_weight =< Allowed of
                true ->
                    New_mean = case Combined_weight of
                        +0.0 -> +0.0;
                        -0.0 -> -0.0;
                        Gleam@denominator@1 -> ((erlang:element(2, Current) * erlang:element(
                            3,
                            Current
                        ))
                        + (erlang:element(2, Next) * erlang:element(3, Next)))
                        / Gleam@denominator@1
                    end,
                    Combined = {centroid, New_mean, Combined_weight},
                    compress_walk(
                        Rest,
                        Combined,
                        Current_q_start,
                        Delta,
                        Total,
                        Acc
                    );

                false ->
                    compress_walk(Rest, Next, Current_q_start + (case Total of
                            +0.0 -> +0.0;
                            -0.0 -> -0.0;
                            Gleam@denominator@2 -> erlang:element(3, Current) / Gleam@denominator@2
                        end), Delta, Total, [Current | Acc])
            end
    end.

-file("src/viva_math/tdigest.gleam", 238).
?DOC(
    " Compress the digest: walk sorted centroids merging adjacent ones while\n"
    " each combined size stays under the k(q)-derived bound. This is the\n"
    " heart of t-digest's accuracy guarantee.\n"
).
-spec compress(t_digest()) -> t_digest().
compress(Td) ->
    Total = erlang:element(4, Td),
    Delta = erlang:element(2, Td),
    case erlang:element(3, Td) of
        [] ->
            Td;

        [First | Rest] ->
            {Merged, _} = compress_walk(Rest, First, +0.0, Delta, Total, []),
            {t_digest, Delta, lists:reverse(Merged), Total}
    end.

-file("src/viva_math/tdigest.gleam", 105).
-spec compression_threshold(t_digest()) -> integer().
compression_threshold(Td) ->
    erlang:trunc(6.0 * erlang:element(2, Td)).

-file("src/viva_math/tdigest.gleam", 213).
-spec insert_sorted(list(centroid()), centroid()) -> list(centroid()).
insert_sorted(Xs, C) ->
    case Xs of
        [] ->
            [C];

        [Head | Rest] ->
            case erlang:element(2, C) =< erlang:element(2, Head) of
                true ->
                    [C, Head | Rest];

                false ->
                    [Head | insert_sorted(Rest, C)]
            end
    end.

-file("src/viva_math/tdigest.gleam", 204).
-spec merge_centroid(t_digest(), centroid()) -> t_digest().
merge_centroid(Td, C) ->
    New_centroids = insert_sorted(erlang:element(3, Td), C),
    {t_digest,
        erlang:element(2, Td),
        New_centroids,
        erlang:element(4, Td) + erlang:element(3, C)}.

-file("src/viva_math/tdigest.gleam", 97).
?DOC(" Insert a weighted sample.\n").
-spec insert_weighted(t_digest(), float(), float()) -> t_digest().
insert_weighted(Td, Value, Weight) ->
    Merged = merge_centroid(Td, {centroid, Value, Weight}),
    case erlang:length(erlang:element(3, Merged)) > compression_threshold(
        Merged
    ) of
        true ->
            compress(Merged);

        false ->
            Merged
    end.

-file("src/viva_math/tdigest.gleam", 92).
?DOC(" Insert a single sample into the digest.\n").
-spec insert(t_digest(), float()) -> t_digest().
insert(Td, Value) ->
    insert_weighted(Td, Value, 1.0).

-file("src/viva_math/tdigest.gleam", 115).
?DOC(" Insert many samples from a list. Equivalent to folding `insert`.\n").
-spec insert_all(t_digest(), list(float())) -> t_digest().
insert_all(Td, Xs) ->
    gleam@list:fold(Xs, Td, fun insert/2).

-file("src/viva_math/tdigest.gleam", 224).
-spec sorted_concat(list(centroid()), list(centroid())) -> list(centroid()).
sorted_concat(A, B) ->
    Merged = lists:append(A, B),
    gleam@list:sort(
        Merged,
        fun(X, Y) ->
            case {erlang:element(2, X) < erlang:element(2, Y),
                erlang:element(2, X) > erlang:element(2, Y)} of
                {true, _} ->
                    lt;

                {_, true} ->
                    gt;

                {_, _} ->
                    eq
            end
        end
    ).

-file("src/viva_math/tdigest.gleam", 302).
-spec float_max(float(), float()) -> float().
float_max(A, B) ->
    case A > B of
        true ->
            A;

        false ->
            B
    end.

-file("src/viva_math/tdigest.gleam", 120).
?DOC(" Combine two digests into one. Useful for distributed reductions.\n").
-spec merge(t_digest(), t_digest()) -> t_digest().
merge(A, B) ->
    Combined = {t_digest,
        float_max(erlang:element(2, A), erlang:element(2, B)),
        sorted_concat(erlang:element(3, A), erlang:element(3, B)),
        erlang:element(4, A) + erlang:element(4, B)},
    compress(Combined).

-file("src/viva_math/tdigest.gleam", 146).
-spec scan_quantile(list(centroid()), float(), float()) -> float().
scan_quantile(Centroids, Target, Acc) ->
    case Centroids of
        [] ->
            +0.0;

        [C] ->
            erlang:element(2, C);

        [A, B | Rest] ->
            Cumulative = Acc + erlang:element(3, A),
            case Target =< Cumulative of
                true ->
                    Span = erlang:element(3, A) + erlang:element(3, B),
                    Position = case Span of
                        +0.0 -> +0.0;
                        -0.0 -> -0.0;
                        Gleam@denominator -> (Target - Acc) / Gleam@denominator
                    end,
                    erlang:element(2, A) + (Position * (erlang:element(2, B) - erlang:element(
                        2,
                        A
                    )));

                false ->
                    scan_quantile([B | Rest], Target, Cumulative)
            end
    end.

-file("src/viva_math/tdigest.gleam", 135).
?DOC(" Approximate quantile q ∈ [0, 1]. Returns `Error` for an empty digest.\n").
-spec quantile(t_digest(), float()) -> {ok, float()} | {error, nil}.
quantile(Td, Q) ->
    case {erlang:element(3, Td), (Q < +0.0) orelse (Q > 1.0)} of
        {[], _} ->
            {error, nil};

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

        {_, false} ->
            Target = Q * erlang:element(4, Td),
            {ok, scan_quantile(erlang:element(3, Td), Target, +0.0)}
    end.

-file("src/viva_math/tdigest.gleam", 170).
?DOC(" Total number of samples represented by the digest.\n").
-spec count(t_digest()) -> float().
count(Td) ->
    erlang:element(4, Td).

-file("src/viva_math/tdigest.gleam", 175).
?DOC(" Minimum sample so far (first centroid mean, since centroids stay sorted).\n").
-spec min(t_digest()) -> {ok, float()} | {error, nil}.
min(Td) ->
    case erlang:element(3, Td) of
        [] ->
            {error, nil};

        [C | _] ->
            {ok, erlang:element(2, C)}
    end.

-file("src/viva_math/tdigest.gleam", 183).
?DOC(" Maximum sample so far.\n").
-spec max(t_digest()) -> {ok, float()} | {error, nil}.
max(Td) ->
    case gleam@list:last(erlang:element(3, Td)) of
        {ok, C} ->
            {ok, erlang:element(2, C)};

        {error, _} ->
            {error, nil}
    end.

-file("src/viva_math/tdigest.gleam", 191).
?DOC(" Convenience: median.\n").
-spec median(t_digest()) -> {ok, float()} | {error, nil}.
median(Td) ->
    quantile(Td, 0.5).

-file("src/viva_math/tdigest.gleam", 196).
?DOC(" Convenience: p99 (extreme-tail quantile, where t-digest excels).\n").
-spec p99(t_digest()) -> {ok, float()} | {error, nil}.
p99(Td) ->
    quantile(Td, 0.99).