gleam_stats

-module(gleam_stats@distributions@negbinomial).
-compile(no_auto_import).

-export([negbinomial_mean/2, negbinomial_variance/2, negbinomial_pmf/3, negbinomial_cdf/3, negbinomial_random/4]).

-spec check_negbinomial_parameters(integer(), float()) -> {ok, boolean()} |
    {error, binary()}.
check_negbinomial_parameters(R, P) ->
    case R > 0 of
        false ->
            _pipe = <<"Invalid input argument: r <= 0. Valid input is r > 0."/utf8>>,
            {error, _pipe};

        true ->
            case (0.0 =< P) andalso (P =< 1.0) of
                false ->
                    _pipe@1 = <<"Invalid input argument: p < 0 or p > 1. Valid input is 0 <= p <= 1."/utf8>>,
                    {error, _pipe@1};

                true ->
                    _pipe@2 = true,
                    {ok, _pipe@2}
            end
    end.

-spec negbinomial_mean(integer(), float()) -> {ok, float()} | {error, binary()}.
negbinomial_mean(R, P) ->
    case check_negbinomial_parameters(R, P) of
        {error, String} ->
            _pipe = String,
            {error, _pipe};

        _@1 ->
            _pipe@1 = case (1.0 - P) of
                0.0 -> 0.0;
                Gleam@denominator -> (gleam@int:to_float(R) * P) / Gleam@denominator
            end,
            {ok, _pipe@1}
    end.

-spec negbinomial_variance(integer(), float()) -> {ok, float()} |
    {error, binary()}.
negbinomial_variance(R, P) ->
    case check_negbinomial_parameters(R, P) of
        {error, String} ->
            _pipe = String,
            {error, _pipe};

        _@1 ->
            {ok, V@1} = case gleam_stats@math:pow(1.0 - P, 2.0) of
                {ok, V} -> {ok, V};
                _try ->
                    erlang:error(#{gleam_error => assert,
                                   message => <<"Assertion pattern match failed"/utf8>>,
                                   value => _try,
                                   module => <<"gleam_stats/distributions/negbinomial"/utf8>>,
                                   function => <<"negbinomial_variance"/utf8>>,
                                   line => 92})
            end,
            _pipe@1 = case V@1 of
                0.0 -> 0.0;
                Gleam@denominator -> (gleam@int:to_float(R) * P) / Gleam@denominator
            end,
            {ok, _pipe@1}
    end.

-spec negbinomial_pmf(integer(), integer(), float()) -> {ok, float()} |
    {error, binary()}.
negbinomial_pmf(X, R, P) ->
    do_negbinomial_pmf(X, R, P).

-spec do_negbinomial_pmf(integer(), integer(), float()) -> {ok, float()} |
    {error, binary()}.
do_negbinomial_pmf(X, R, P) ->
    case check_negbinomial_parameters(R, P) of
        {error, String} ->
            _pipe = String,
            {error, _pipe};

        _@1 ->
            case (X >= 0) andalso (((X + R) - 1) > 0) of
                true ->
                    {ok, C@1} = case gleam_stats@math:combination(
                        (X
                        + R)
                        - 1,
                        X
                    ) of
                        {ok, C} -> {ok, C};
                        _try ->
                            erlang:error(#{gleam_error => assert,
                                           message => <<"Assertion pattern match failed"/utf8>>,
                                           value => _try,
                                           module => <<"gleam_stats/distributions/negbinomial"/utf8>>,
                                           function => <<"do_negbinomial_pmf"/utf8>>,
                                           line => 153})
                    end,
                    {ok, V1@1} = case gleam_stats@math:pow(
                        1.0
                        - P,
                        gleam@int:to_float(R)
                    ) of
                        {ok, V1} -> {ok, V1};
                        _try@1 ->
                            erlang:error(#{gleam_error => assert,
                                           message => <<"Assertion pattern match failed"/utf8>>,
                                           value => _try@1,
                                           module => <<"gleam_stats/distributions/negbinomial"/utf8>>,
                                           function => <<"do_negbinomial_pmf"/utf8>>,
                                           line => 154})
                    end,
                    {ok, V2@1} = case gleam_stats@math:pow(
                        P,
                        gleam@int:to_float(X)
                    ) of
                        {ok, V2} -> {ok, V2};
                        _try@2 ->
                            erlang:error(#{gleam_error => assert,
                                           message => <<"Assertion pattern match failed"/utf8>>,
                                           value => _try@2,
                                           module => <<"gleam_stats/distributions/negbinomial"/utf8>>,
                                           function => <<"do_negbinomial_pmf"/utf8>>,
                                           line => 155})
                    end,
                    _pipe@1 = (gleam@int:to_float(C@1) * V1@1) * V2@1,
                    {ok, _pipe@1};

                _@2 ->
                    _pipe@2 = 0.0,
                    {ok, _pipe@2}
            end
    end.

-spec negbinomial_cdf(integer(), integer(), float()) -> {ok, float()} |
    {error, binary()}.
negbinomial_cdf(X, R, P) ->
    do_negbinomial_cdf(X, R, P).

-spec do_negbinomial_cdf(integer(), integer(), float()) -> {ok, float()} |
    {error, binary()}.
do_negbinomial_cdf(X, R, P) ->
    case check_negbinomial_parameters(R, P) of
        {error, String} ->
            _pipe = String,
            {error, _pipe};

        _@1 ->
            case X < 0 of
                true ->
                    _pipe@1 = 0.0,
                    {ok, _pipe@1};

                false ->
                    case (X >= 0) andalso (((X + R) - 1) > 0) of
                        true ->
                            _pipe@2 = gleam@list:range(0, X + 1),
                            _pipe@3 = gleam@list:fold(
                                _pipe@2,
                                0.0,
                                fun(Acc, I) ->
                                    {ok, C@1} = case gleam_stats@math:combination(
                                        (I
                                        + R)
                                        - 1,
                                        I
                                    ) of
                                        {ok, C} -> {ok, C};
                                        _try ->
                                            erlang:error(
                                                    #{gleam_error => assert,
                                                      message => <<"Assertion pattern match failed"/utf8>>,
                                                      value => _try,
                                                      module => <<"gleam_stats/distributions/negbinomial"/utf8>>,
                                                      function => <<"do_negbinomial_cdf"/utf8>>,
                                                      line => 236}
                                                )
                                    end,
                                    {ok, V1@1} = case gleam_stats@math:pow(
                                        1.0
                                        - P,
                                        gleam@int:to_float(R)
                                    ) of
                                        {ok, V1} -> {ok, V1};
                                        _try@1 ->
                                            erlang:error(
                                                    #{gleam_error => assert,
                                                      message => <<"Assertion pattern match failed"/utf8>>,
                                                      value => _try@1,
                                                      module => <<"gleam_stats/distributions/negbinomial"/utf8>>,
                                                      function => <<"do_negbinomial_cdf"/utf8>>,
                                                      line => 237}
                                                )
                                    end,
                                    {ok, V2@1} = case gleam_stats@math:pow(
                                        P,
                                        gleam@int:to_float(I)
                                    ) of
                                        {ok, V2} -> {ok, V2};
                                        _try@2 ->
                                            erlang:error(
                                                    #{gleam_error => assert,
                                                      message => <<"Assertion pattern match failed"/utf8>>,
                                                      value => _try@2,
                                                      module => <<"gleam_stats/distributions/negbinomial"/utf8>>,
                                                      function => <<"do_negbinomial_cdf"/utf8>>,
                                                      line => 238}
                                                )
                                    end,
                                    Acc
                                    + ((gleam@int:to_float(C@1)
                                    * V1@1)
                                    * V2@1)
                                end
                            ),
                            {ok, _pipe@3};

                        false ->
                            _pipe@4 = 1.0,
                            {ok, _pipe@4}
                    end
            end
    end.

-spec negbinomial_random(
    gleam@iterator:iterator(integer()),
    integer(),
    float(),
    integer()
) -> {ok, {list(integer()), gleam@iterator:iterator(integer())}} |
    {error, binary()}.
negbinomial_random(Stream, R, P, M) ->
    do_negbinomial_random(Stream, R, P, M).

-spec do_negbinomial_random(
    gleam@iterator:iterator(integer()),
    integer(),
    float(),
    integer()
) -> {ok, {list(integer()), gleam@iterator:iterator(integer())}} |
    {error, binary()}.
do_negbinomial_random(Stream, R, P, M) ->
    case check_negbinomial_parameters(R, P) of
        {error, String} ->
            _pipe = String,
            {error, _pipe};

        _@1 ->
            case M > 0 of
                false ->
                    _pipe@1 = <<"Invalid input arugment: m < 0. Valid input is m > 0."/utf8>>,
                    {error, _pipe@1};

                true ->
                    {ok, Out@1} = case gleam_stats@distributions@geometric:geometric_random(
                        Stream,
                        P,
                        R
                        * M
                    ) of
                        {ok, Out} -> {ok, Out};
                        _try ->
                            erlang:error(#{gleam_error => assert,
                                           message => <<"Assertion pattern match failed"/utf8>>,
                                           value => _try,
                                           module => <<"gleam_stats/distributions/negbinomial"/utf8>>,
                                           function => <<"do_negbinomial_random"/utf8>>,
                                           line => 323})
                    end,
                    Numbers = begin
                        _pipe@2 = gleam@pair:first(Out@1),
                        _pipe@3 = gleam@list:window(_pipe@2, R),
                        gleam@list:map(
                            _pipe@3,
                            fun(X) ->
                                _pipe@4 = X,
                                gleam@list:fold(
                                    _pipe@4,
                                    0,
                                    fun(A, B) -> A + B end
                                )
                            end
                        )
                    end,
                    _pipe@5 = {Numbers, gleam@pair:second(Out@1)},
                    {ok, _pipe@5}
            end
    end.