gleam_stats

-module(gleam_stats@rand).
-compile(no_auto_import).

-export([take_randints/2, next_uniform/4, next_normal/4, next_randint/4, next_bern/3, next_binom/4, next_negbinom/4, next_geom/3, next_exp/3]).

-spec box_muller(float(), float()) -> float().
box_muller(U1, U2) ->
    {ok, X@1} = case gleam@float:square_root(-2.0 * math:log(U1)) of
        {ok, X} -> {ok, X};
        _try ->
            erlang:error(#{gleam_error => assert,
                           message => <<"Assertion pattern match failed"/utf8>>,
                           value => _try,
                           module => <<"gleam_stats/rand"/utf8>>,
                           function => <<"box_muller"/utf8>>,
                           line => 48})
    end,
    X@1 * math:cos((2.0 * math:pi()) * U2).

-spec take_randints(gleam@iterator:iterator(integer()), integer()) -> {list(integer()),
                                                                       gleam@iterator:iterator(integer())}.
take_randints(Stream, M) ->
    Numbers = begin
        _pipe = Stream,
        _pipe@1 = gleam@iterator:take(_pipe, M),
        gleam@iterator:to_list(_pipe@1)
    end,
    {Numbers, gleam@iterator:drop(Stream, M)}.

-spec next_uniform(
    gleam@iterator:iterator(integer()),
    float(),
    float(),
    integer()
) -> {list(float()), gleam@iterator:iterator(integer())}.
next_uniform(Stream, Min, Max, M) ->
    Out = take_randints(Stream, M),
    Numbers = begin
        _pipe = gleam@pair:first(Out),
        gleam@list:map(
            _pipe,
            fun(X) -> (Max * (case gleam@int:to_float(4294967295) of
                    0.0 -> 0.0;
                    Gleam@denominator -> gleam@int:to_float(X) / Gleam@denominator
                end)) + Min end
        )
    end,
    {Numbers, gleam@pair:second(Out)}.

-spec next_normal(
    gleam@iterator:iterator(integer()),
    float(),
    float(),
    integer()
) -> {list(float()), gleam@iterator:iterator(integer())}.
next_normal(Stream, Mean, Sigma, M) ->
    Out = next_uniform(Stream, 0.0, 1.0, 2 * M),
    Numbers = begin
        _pipe = gleam@pair:first(Out),
        _pipe@1 = gleam@list:window(_pipe, 2),
        gleam@list:map(_pipe@1, fun(X) -> case X of
                    [U1, U2] ->
                        (Sigma * box_muller(U1, U2)) + Mean
                end end)
    end,
    {Numbers, gleam@pair:second(Out)}.

-spec next_randint(
    gleam@iterator:iterator(integer()),
    integer(),
    integer(),
    integer()
) -> {list(integer()), gleam@iterator:iterator(integer())}.
next_randint(Stream, Min, Max, M) ->
    C = (gleam@int:to_float(Max) - gleam@int:to_float(Min)) + 1.0,
    Out = next_uniform(Stream, 0.0, 1.0, M),
    Numbers = begin
        _pipe = gleam@pair:first(Out),
        gleam@list:map(
            _pipe,
            fun(X) ->
                gleam@float:round(
                    gleam@int:to_float(Min)
                    + gleam@float:floor(C * X)
                )
            end
        )
    end,
    {Numbers, gleam@pair:second(Out)}.

-spec next_bern(gleam@iterator:iterator(integer()), float(), integer()) -> {list(integer()),
                                                                            gleam@iterator:iterator(integer())}.
next_bern(Stream, P, M) ->
    Out = next_uniform(Stream, 0.0, 1.0, M),
    Numbers = begin
        _pipe = gleam@pair:first(Out),
        gleam@list:map(_pipe, fun(X) -> case X =< P of
                    true ->
                        1;

                    false ->
                        0
                end end)
    end,
    {Numbers, gleam@pair:second(Out)}.

-spec next_binom(
    gleam@iterator:iterator(integer()),
    float(),
    integer(),
    integer()
) -> {list(integer()), gleam@iterator:iterator(integer())}.
next_binom(Stream, P, N, M) ->
    Out = next_bern(Stream, P, N * M),
    Numbers = begin
        _pipe = gleam@pair:first(Out),
        _pipe@1 = gleam@list:window(_pipe, N),
        gleam@list:map(
            _pipe@1,
            fun(X) ->
                _pipe@2 = X,
                gleam@list:fold(_pipe@2, 0, fun(A, B) -> A + B end)
            end
        )
    end,
    {Numbers, gleam@pair:second(Out)}.

-spec next_negbinom(
    gleam@iterator:iterator(integer()),
    float(),
    integer(),
    integer()
) -> {list(integer()), gleam@iterator:iterator(integer())}.
next_negbinom(Stream, P, N, M) ->
    Out = next_geom(Stream, P, N * M),
    Numbers = begin
        _pipe = gleam@pair:first(Out),
        _pipe@1 = gleam@list:window(_pipe, N),
        gleam@list:map(
            _pipe@1,
            fun(X) ->
                _pipe@2 = X,
                gleam@list:fold(_pipe@2, 0, fun(A, B) -> A + B end)
            end
        )
    end,
    {Numbers, gleam@pair:second(Out)}.

-spec next_geom(gleam@iterator:iterator(integer()), float(), integer()) -> {list(integer()),
                                                                            gleam@iterator:iterator(integer())}.
next_geom(Stream, P, M) ->
    Out = next_uniform(Stream, 0.0, 1.0, M),
    Numbers = begin
        _pipe = gleam@pair:first(Out),
        gleam@list:map(
            _pipe,
            fun(X) ->
                gleam@float:round(gleam@float:floor(case math:log(1.0 - P) of
                            0.0 -> 0.0;
                            Gleam@denominator -> math:log(X) / Gleam@denominator
                        end))
            end
        )
    end,
    {Numbers, gleam@pair:second(Out)}.

-spec next_exp(gleam@iterator:iterator(integer()), float(), integer()) -> {list(float()),
                                                                           gleam@iterator:iterator(integer())}.
next_exp(Stream, Lambda, M) ->
    Out = next_uniform(Stream, 0.0, 1.0, M),
    Numbers = begin
        _pipe = gleam@pair:first(Out),
        gleam@list:map(_pipe, fun(X) -> (case (-1.0 * Lambda) of
                    0.0 -> 0.0;
                    Gleam@denominator -> 1.0 / Gleam@denominator
                end) * math:log(X) end)
    end,
    {Numbers, gleam@pair:second(Out)}.