gleam_stats

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

-export([seed_pcg32/2, seed_lcg32/1, take_randints/2, uint32/1, uint64/1]).
-export_type([permuted_congruential_generator/0, linear_congruential_generator/0]).

-type permuted_congruential_generator() :: {permuted_congruential_generator,
                                            integer(),
                                            integer(),
                                            integer(),
                                            integer(),
                                            integer(),
                                            integer()}.

-type linear_congruential_generator() :: {linear_congruential_generator,
                                          integer(),
                                          integer()}.

-spec pcg32_next_rn({integer(), integer()}, permuted_congruential_generator()) -> integer().
pcg32_next_rn(State, Pcg) ->
    Old_state = gleam@pair:first(State),
    Xorshifted = uint32(
        gleam@bitwise:shift_right(
            gleam@bitwise:exclusive_or(
                gleam@bitwise:shift_right(Old_state, erlang:element(3, Pcg)),
                Old_state
            ),
            erlang:element(4, Pcg)
        )
    ),
    Rotation = uint32(
        gleam@bitwise:shift_right(Old_state, erlang:element(5, Pcg))
    ),
    uint32(
        gleam@bitwise:'or'(
            gleam@bitwise:shift_right(Xorshifted, Rotation),
            gleam@bitwise:shift_left(
                Xorshifted,
                gleam@bitwise:'and'(-1 * Rotation, erlang:element(6, Pcg))
            )
        )
    ).

-spec pcg32_init(integer(), integer(), permuted_congruential_generator()) -> {integer(),
                                                                              integer()}.
pcg32_init(Seed, Seq, Pcg) ->
    _pipe = {0,
             gleam@bitwise:'or'(
                 uint64(
                     gleam@bitwise:shift_left(
                         Seq,
                         erlang:element(
                             2,
                             {permuted_congruential_generator,
                              1,
                              18,
                              27,
                              59,
                              31,
                              6364136223846793005}
                         )
                     )
                 ),
                 erlang:element(2, Pcg)
             )},
    _pipe@1 = pcg32_next_state(_pipe, Pcg),
    _pipe@2 = (fun(State) ->
        {S, I} = State,
        {uint64(S + Seed), I}
    end)(_pipe@1),
    pcg32_next_state(_pipe@2, Pcg).

-spec pcg32_next_state(
    {integer(), integer()},
    permuted_congruential_generator()
) -> {integer(), integer()}.
pcg32_next_state(State, Pcg) ->
    {S, I} = State,
    {uint64((S * erlang:element(7, Pcg)) + I), I}.

-spec seed_pcg32(integer(), integer()) -> gleam@iterator:iterator(integer()).
seed_pcg32(Seed, Seq) ->
    Pcg = {permuted_congruential_generator,
           1,
           18,
           27,
           59,
           31,
           6364136223846793005},
    _pipe = pcg32_init(uint64(Seed), uint64(Seq), Pcg),
    gleam@iterator:unfold(
        _pipe,
        fun(State) ->
            Next_rn = pcg32_next_rn(State, Pcg),
            Next_state = pcg32_next_state(State, Pcg),
            {next, Next_rn, Next_state}
        end
    ).

-spec lcg32_init(integer(), linear_congruential_generator()) -> integer().
lcg32_init(Seed, Lcg) ->
    gleam@bitwise:'and'(
        (erlang:element(2, Lcg)
        * Seed)
        + erlang:element(3, Lcg),
        4294967295
    ).

-spec lcg32_next_state(integer(), linear_congruential_generator()) -> integer().
lcg32_next_state(State, Lcg) ->
    gleam@bitwise:'and'(
        (erlang:element(2, Lcg)
        * State)
        + erlang:element(3, Lcg),
        4294967295
    ).

-spec seed_lcg32(integer()) -> gleam@iterator:iterator(integer()).
seed_lcg32(Seed) ->
    Lcg = {linear_congruential_generator, 1664525, 1013904223},
    _pipe = lcg32_init(Seed, Lcg),
    gleam@iterator:unfold(
        _pipe,
        fun(State) ->
            Next_state = lcg32_next_state(State, Lcg),
            {next, Next_state, Next_state}
        end
    ).

-spec take_randints(gleam@iterator:iterator(integer()), integer()) -> {ok,
                                                                       {list(integer()),
                                                                        gleam@iterator:iterator(integer())}} |
    {error, binary()}.
take_randints(Stream, M) ->
    case M > 0 of
        false ->
            {error,
             <<"Invalid input arugment: m < 0. Valid input is m > 0."/utf8>>};

        true ->
            Numbers = begin
                _pipe = Stream,
                _pipe@1 = gleam@iterator:take(_pipe, M),
                gleam@iterator:to_list(_pipe@1)
            end,
            _pipe@2 = {Numbers, gleam@iterator:drop(Stream, M)},
            {ok, _pipe@2}
    end.

-spec uint32(integer()) -> integer().
uint32(N) ->
    gleam@bitwise:'and'(N, 4294967295).

-spec uint64(integer()) -> integer().
uint64(N) ->
    gleam@bitwise:'and'(N, 18446744073709551615).