viva_math

-module(viva_math@random).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_math/random.gleam").
-export([from_int/1, with_algo/2, jump/1, uniform/1, uniform_in/3, integer/2, standard_normal/1, normal/3, bernoulli/2, categorical/2, choice/2, shuffle/2, standard_normals/2, uniforms/2]).
-export_type([seed/0, algorithm/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(
    " Pure functional random number generation.\n"
    "\n"
    " Built on top of Erlang's `:rand` module (OTP 22+). The `Seed` type is\n"
    " **opaque and immutable**: every sampling function returns a new seed,\n"
    " allowing fully reproducible deterministic streams without process state.\n"
    "\n"
    " The default algorithm is `exsss` (Xorshift116** with StarStar scrambler,\n"
    " period 2¹¹⁶-1, 58-bit output). For longer streams or jump-ahead, use\n"
    " `with_algo(Exro928ss, seed)`.\n"
    "\n"
    " ## Why not the hash-based generator in `common`?\n"
    "\n"
    " `common.deterministic_noise` is fine for reproducible *fixtures*, but it\n"
    " is statistically weak. For Monte Carlo, neural net initialization or\n"
    " Bayesian sampling, prefer this module.\n"
    "\n"
    " ## Example\n"
    "\n"
    " ```gleam\n"
    " import viva_math/random\n"
    "\n"
    " let seed0 = random.from_int(42)\n"
    " let #(x, seed1) = random.uniform(seed0)\n"
    " let #(y, seed2) = random.normal(seed1, 0.0, 1.0)\n"
    " // Same seed0 always produces the same x, y sequence.\n"
    " ```\n"
).

-type seed() :: any().

-type algorithm() :: exsss | exro928ss | exrop | exsp | exs1024s | mwc59.

-file("src/viva_math/random.gleam", 94).
?DOC(" Build a seed from an integer using the default algorithm (Exsss).\n").
-spec from_int(integer()) -> seed().
from_int(Seed) ->
    viva_math_random_ffi:seed_default(Seed).

-file("src/viva_math/random.gleam", 99).
?DOC(" Build a seed with an explicit algorithm.\n").
-spec with_algo(algorithm(), integer()) -> seed().
with_algo(Algorithm, Seed) ->
    viva_math_random_ffi:seed_with_algo(Algorithm, Seed).

-file("src/viva_math/random.gleam", 104).
?DOC(" Advance the seed by 2⁶⁴ calls. Useful for non-overlapping parallel streams.\n").
-spec jump(seed()) -> seed().
jump(Seed) ->
    viva_math_random_ffi:jump(Seed).

-file("src/viva_math/random.gleam", 113).
?DOC(" Uniform float in [0.0, 1.0).\n").
-spec uniform(seed()) -> {float(), seed()}.
uniform(Seed) ->
    viva_math_random_ffi:uniform_real(Seed).

-file("src/viva_math/random.gleam", 118).
?DOC(" Uniform float in [low, high).\n").
-spec uniform_in(seed(), float(), float()) -> {float(), seed()}.
uniform_in(Seed, Low, High) ->
    {U, S} = viva_math_random_ffi:uniform_real(Seed),
    {Low + (U * (High - Low)), S}.

-file("src/viva_math/random.gleam", 124).
?DOC(" Uniform integer in [1, n]. Returns an error if n < 1.\n").
-spec integer(seed(), integer()) -> {ok, {integer(), seed()}} | {error, nil}.
integer(Seed, N) ->
    case N < 1 of
        true ->
            {error, nil};

        false ->
            {ok, viva_math_random_ffi:uniform_int(N, Seed)}
    end.

-file("src/viva_math/random.gleam", 132).
?DOC(" Standard normal N(0, 1).\n").
-spec standard_normal(seed()) -> {float(), seed()}.
standard_normal(Seed) ->
    viva_math_random_ffi:normal_standard(Seed).

-file("src/viva_math/random.gleam", 137).
?DOC(" Normal N(mu, sigma²).\n").
-spec normal(seed(), float(), float()) -> {float(), seed()}.
normal(Seed, Mu, Sigma) ->
    viva_math_random_ffi:normal_with(Mu, Sigma, Seed).

-file("src/viva_math/random.gleam", 142).
?DOC(" Bernoulli sample: True with probability `p`. `p` is clamped to [0, 1].\n").
-spec bernoulli(seed(), float()) -> {boolean(), seed()}.
bernoulli(Seed, P) ->
    P_clamped = case P of
        P@1 when P@1 < +0.0 ->
            +0.0;

        P@2 when P@2 > 1.0 ->
            1.0;

        _ ->
            P
    end,
    {U, S} = viva_math_random_ffi:uniform_real(Seed),
    {U < P_clamped, S}.

-file("src/viva_math/random.gleam", 178).
-spec categorical_walk(list(float()), float(), integer(), float()) -> integer().
categorical_walk(Probs, Target, Index, Acc) ->
    case Probs of
        [] ->
            Index - 1;

        [P | Rest] ->
            Next_acc = Acc + P,
            case Target < Next_acc of
                true ->
                    Index;

                false ->
                    categorical_walk(Rest, Target, Index + 1, Next_acc)
            end
    end.

-file("src/viva_math/random.gleam", 160).
?DOC(
    " Sample an index from a categorical distribution defined by `probs`.\n"
    "\n"
    " Probabilities are renormalised internally. Returns an error if the list\n"
    " is empty or sums to a non-positive value.\n"
).
-spec categorical(seed(), list(float())) -> {ok, {integer(), seed()}} |
    {error, nil}.
categorical(Seed, Probs) ->
    Any_negative = gleam@list:any(Probs, fun(P) -> P < +0.0 end),
    Total = gleam@list:fold(Probs, +0.0, fun(Acc, P@1) -> Acc + P@1 end),
    case {Probs, Any_negative, Total > +0.0} of
        {[], _, _} ->
            {error, nil};

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

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

        {_, _, true} ->
            {U, S} = viva_math_random_ffi:uniform_real(Seed),
            Target = U * Total,
            {ok, {categorical_walk(Probs, Target, 0, +0.0), S}}
    end.

-file("src/viva_math/random.gleam", 306).
-spec list_at(list(EXY), integer()) -> {ok, EXY} | {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/random.gleam", 197).
?DOC(" Pick a uniformly random element from a list.\n").
-spec choice(seed(), list(EWS)) -> {ok, {EWS, seed()}} | {error, nil}.
choice(Seed, Xs) ->
    case erlang:length(Xs) of
        0 ->
            {error, nil};

        N ->
            {Idx, S} = viva_math_random_ffi:uniform_int(N, Seed),
            case list_at(Xs, Idx - 1) of
                {ok, V} ->
                    {ok, {V, S}};

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

-file("src/viva_math/random.gleam", 238).
-spec dict_to_list(
    gleam@dict:dict(integer(), EXF),
    integer(),
    integer(),
    list(EXF)
) -> list(EXF).
dict_to_list(D, I, N, Acc) ->
    case I >= N of
        true ->
            lists:reverse(Acc);

        false ->
            case gleam_stdlib:map_get(D, I) of
                {ok, V} ->
                    dict_to_list(D, I + 1, N, [V | Acc]);

                {error, _} ->
                    lists:reverse(Acc)
            end
    end.

-file("src/viva_math/random.gleam", 266).
-spec swap(gleam@dict:dict(integer(), EXP), integer(), integer()) -> gleam@dict:dict(integer(), EXP).
swap(D, I, J) ->
    case {gleam_stdlib:map_get(D, I), gleam_stdlib:map_get(D, J)} of
        {{ok, Vi}, {ok, Vj}} ->
            gleam@dict:insert(gleam@dict:insert(D, I, Vj), J, Vi);

        {_, _} ->
            D
    end.

-file("src/viva_math/random.gleam", 249).
-spec fisher_yates(gleam@dict:dict(integer(), EXK), integer(), seed()) -> {gleam@dict:dict(integer(), EXK),
    seed()}.
fisher_yates(Arr, I, Seed) ->
    case I =< 0 of
        true ->
            {Arr, Seed};

        false ->
            {J_raw, S1} = viva_math_random_ffi:uniform_int(I + 1, Seed),
            J = J_raw - 1,
            Swapped = swap(Arr, I, J),
            fisher_yates(Swapped, I - 1, S1)
    end.

-file("src/viva_math/random.gleam", 231).
-spec list_to_dict(list(EWZ), integer(), gleam@dict:dict(integer(), EWZ)) -> gleam@dict:dict(integer(), EWZ).
list_to_dict(Xs, I, Acc) ->
    case Xs of
        [] ->
            Acc;

        [X | Rest] ->
            list_to_dict(Rest, I + 1, gleam@dict:insert(Acc, I, X))
    end.

-file("src/viva_math/random.gleam", 219).
?DOC(
    " Return a random permutation of `xs` via Fisher–Yates (modern Durstenfeld\n"
    " variant), implemented on top of a `dict` for O(log n) indexed swaps.\n"
    "\n"
    " Truly uniform over all n! permutations (no float-key collision bias the\n"
    " sort-based shuffle suffered from).\n"
).
-spec shuffle(seed(), list(EWW)) -> {list(EWW), seed()}.
shuffle(Seed, Xs) ->
    N = erlang:length(Xs),
    case N of
        0 ->
            {[], Seed};

        _ ->
            Initial = list_to_dict(Xs, 0, maps:new()),
            {Shuffled, Final_seed} = fisher_yates(Initial, N - 1, Seed),
            {dict_to_list(Shuffled, 0, N, []), Final_seed}
    end.

-file("src/viva_math/random.gleam", 287).
-spec draw(integer(), seed(), fun((seed()) -> {float(), seed()}), list(float())) -> {list(float()),
    seed()}.
draw(N, Seed, Sampler, Acc) ->
    case N =< 0 of
        true ->
            {lists:reverse(Acc), Seed};

        false ->
            {X, S} = Sampler(Seed),
            draw(N - 1, S, Sampler, [X | Acc])
    end.

-file("src/viva_math/random.gleam", 278).
?DOC(" Sample n standard normals at once.\n").
-spec standard_normals(seed(), integer()) -> {list(float()), seed()}.
standard_normals(Seed, N) ->
    draw(N, Seed, fun viva_math_random_ffi:normal_standard/1, []).

-file("src/viva_math/random.gleam", 283).
?DOC(" Sample n uniforms in [0, 1) at once.\n").
-spec uniforms(seed(), integer()) -> {list(float()), seed()}.
uniforms(Seed, N) ->
    draw(N, Seed, fun viva_math_random_ffi:uniform_real/1, []).