-module(mlx_random).

%% Random number generation functions for MLX
-export([
    %% Basic random operations
    seed/1, key/1,
    normal/1, normal/2, normal/3,
    uniform/1, uniform/2, uniform/3, uniform/4,
    randint/2, randint/3, randint/4,
    
    %% Statistical distributions
    bernoulli/1, bernoulli/2, bernoulli/3,
    categorical/1, categorical/2, categorical/3,
    multinomial/2, multinomial/3, multinomial/4,
    
    %% Advanced distributions
    gamma/2, gamma/3, gamma/4,
    beta/2, beta/3, beta/4,
    exponential/1, exponential/2, exponential/3,
    poisson/1, poisson/2, poisson/3,
    
    %% Random utilities
    shuffle/1, shuffle/2,
    choice/1, choice/2, choice/3,
    permutation/1, permutation/2
]).

%% Type definitions
-type array() :: reference().
-type shape() :: [integer()].
-type dtype() :: atom().
-type key() :: reference().

%% Random seed management
-spec seed(integer()) -> ok | {error, term()}.
seed(Seed) ->
    mlx_random_nif:seed(Seed).

-spec key(integer()) -> {ok, key()} | {error, term()}.
key(Seed) ->
    mlx_random_nif:key(Seed).

%% Normal distribution
-spec normal(shape()) -> {ok, array()} | {error, term()}.
normal(Shape) ->
    normal(Shape, 0.0, 1.0).

-spec normal(shape(), number()) -> {ok, array()} | {error, term()}.
normal(Shape, Mean) ->
    normal(Shape, Mean, 1.0).

-spec normal(shape(), number(), number()) -> {ok, array()} | {error, term()}.
normal(Shape, Mean, Std) ->
    mlx_random_nif:normal(Shape, Mean, Std).

%% Uniform distribution
-spec uniform(shape()) -> {ok, array()} | {error, term()}.
uniform(Shape) ->
    uniform(Shape, 0.0, 1.0).

-spec uniform(shape(), number()) -> {ok, array()} | {error, term()}.
uniform(Shape, Low) ->
    uniform(Shape, Low, 1.0).

-spec uniform(shape(), number(), number()) -> {ok, array()} | {error, term()}.
uniform(Shape, Low, High) ->
    uniform(Shape, Low, High, float32).

-spec uniform(shape(), number(), number(), dtype()) -> {ok, array()} | {error, term()}.
uniform(Shape, Low, High, Dtype) ->
    mlx_random_nif:uniform(Shape, Low, High, Dtype).

%% Random integers
-spec randint(integer(), shape()) -> {ok, array()} | {error, term()}.
randint(High, Shape) ->
    randint(0, High, Shape).

-spec randint(integer(), integer(), shape()) -> {ok, array()} | {error, term()}.
randint(Low, High, Shape) ->
    randint(Low, High, Shape, int32).

-spec randint(integer(), integer(), shape(), dtype()) -> {ok, array()} | {error, term()}.
randint(Low, High, Shape, Dtype) ->
    mlx_random_nif:randint(Low, High, Shape, Dtype).

%% Bernoulli distribution
-spec bernoulli(array()) -> {ok, array()} | {error, term()}.
bernoulli(P) ->
    bernoulli(P, []).

-spec bernoulli(array(), shape()) -> {ok, array()} | {error, term()}.
bernoulli(P, Shape) ->
    bernoulli(P, Shape, float32).

-spec bernoulli(array(), shape(), dtype()) -> {ok, array()} | {error, term()}.
bernoulli(P, Shape, Dtype) ->
    mlx_random_nif:bernoulli(P, Shape, Dtype).

%% Categorical distribution
-spec categorical(array()) -> {ok, array()} | {error, term()}.
categorical(Logits) ->
    categorical(Logits, []).

-spec categorical(array(), shape()) -> {ok, array()} | {error, term()}.
categorical(Logits, Shape) ->
    categorical(Logits, Shape, int32).

-spec categorical(array(), shape(), dtype()) -> {ok, array()} | {error, term()}.
categorical(Logits, Shape, Dtype) ->
    mlx_random_nif:categorical(Logits, Shape, Dtype).

%% Multinomial distribution
-spec multinomial(array(), integer()) -> {ok, array()} | {error, term()}.
multinomial(P, N) ->
    multinomial(P, N, []).

-spec multinomial(array(), integer(), shape()) -> {ok, array()} | {error, term()}.
multinomial(P, N, Shape) ->
    multinomial(P, N, Shape, int32).

-spec multinomial(array(), integer(), shape(), dtype()) -> {ok, array()} | {error, term()}.
multinomial(P, N, Shape, Dtype) ->
    mlx_random_nif:multinomial(P, N, Shape, Dtype).

%% Gamma distribution
-spec gamma(array(), array()) -> {ok, array()} | {error, term()}.
gamma(Alpha, Beta) ->
    gamma(Alpha, Beta, []).

-spec gamma(array(), array(), shape()) -> {ok, array()} | {error, term()}.
gamma(Alpha, Beta, Shape) ->
    gamma(Alpha, Beta, Shape, float32).

-spec gamma(array(), array(), shape(), dtype()) -> {ok, array()} | {error, term()}.
gamma(Alpha, Beta, Shape, Dtype) ->
    mlx_random_nif:gamma(Alpha, Beta, Shape, Dtype).

%% Beta distribution
-spec beta(array(), array()) -> {ok, array()} | {error, term()}.
beta(Alpha, Beta) ->
    beta(Alpha, Beta, []).

-spec beta(array(), array(), shape()) -> {ok, array()} | {error, term()}.
beta(Alpha, Beta, Shape) ->
    beta(Alpha, Beta, Shape, float32).

-spec beta(array(), array(), shape(), dtype()) -> {ok, array()} | {error, term()}.
beta(Alpha, Beta, Shape, Dtype) ->
    mlx_random_nif:beta(Alpha, Beta, Shape, Dtype).

%% Exponential distribution
-spec exponential(array()) -> {ok, array()} | {error, term()}.
exponential(Lambda) ->
    exponential(Lambda, []).

-spec exponential(array(), shape()) -> {ok, array()} | {error, term()}.
exponential(Lambda, Shape) ->
    exponential(Lambda, Shape, float32).

-spec exponential(array(), shape(), dtype()) -> {ok, array()} | {error, term()}.
exponential(Lambda, Shape, Dtype) ->
    mlx_random_nif:exponential(Lambda, Shape, Dtype).

%% Poisson distribution
-spec poisson(array()) -> {ok, array()} | {error, term()}.
poisson(Lambda) ->
    poisson(Lambda, []).

-spec poisson(array(), shape()) -> {ok, array()} | {error, term()}.
poisson(Lambda, Shape) ->
    poisson(Lambda, Shape, int32).

-spec poisson(array(), shape(), dtype()) -> {ok, array()} | {error, term()}.
poisson(Lambda, Shape, Dtype) ->
    mlx_random_nif:poisson(Lambda, Shape, Dtype).

%% Random utilities
-spec shuffle(array()) -> {ok, array()} | {error, term()}.
shuffle(Array) ->
    shuffle(Array, 0).

-spec shuffle(array(), integer()) -> {ok, array()} | {error, term()}.
shuffle(Array, Axis) ->
    mlx_random_nif:shuffle(Array, Axis).

-spec choice(array()) -> {ok, array()} | {error, term()}.
choice(Array) ->
    choice(Array, 1).

-spec choice(array(), integer()) -> {ok, array()} | {error, term()}.
choice(Array, Size) ->
    choice(Array, Size, true).

-spec choice(array(), integer(), boolean()) -> {ok, array()} | {error, term()}.
choice(Array, Size, Replace) ->
    mlx_random_nif:choice(Array, Size, Replace).

-spec permutation(integer()) -> {ok, array()} | {error, term()}.
permutation(N) ->
    mlx_random_nif:permutation(N).

-spec permutation(array(), integer()) -> {ok, array()} | {error, term()}.
permutation(Array, Axis) ->
    mlx_random_nif:permutation(Array, Axis).