-module(mlx_random_nif).

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

%% Load the NIF library
-on_load(init/0).

-spec init() -> ok | {error, term()}.
init() ->
    PrivDir = case code:priv_dir(mlx) of
        {error, _} ->
            EbinDir = filename:dirname(code:which(?MODULE)),
            AppPath = filename:dirname(EbinDir),
            filename:join(AppPath, "priv");
        Path ->
            Path
    end,
    erlang:load_nif(filename:join(PrivDir, "mlx_random_nif"), 0).

%% NIF stubs - these will be replaced by the actual NIF implementations
-spec seed(integer()) -> ok | {error, term()}.
seed(_Seed) ->
    erlang:nif_error({nif_not_loaded, ?MODULE}).

-spec key(integer()) -> {ok, reference()} | {error, term()}.
key(_Seed) ->
    erlang:nif_error({nif_not_loaded, ?MODULE}).

-spec normal([integer()], number(), number()) -> {ok, reference()} | {error, term()}.
normal(_Shape, _Mean, _Std) ->
    erlang:nif_error({nif_not_loaded, ?MODULE}).

-spec uniform([integer()], number(), number(), atom()) -> {ok, reference()} | {error, term()}.
uniform(_Shape, _Low, _High, _Dtype) ->
    erlang:nif_error({nif_not_loaded, ?MODULE}).

-spec randint(integer(), integer(), [integer()], atom()) -> {ok, reference()} | {error, term()}.
randint(_Low, _High, _Shape, _Dtype) ->
    erlang:nif_error({nif_not_loaded, ?MODULE}).

-spec bernoulli(reference(), [integer()], atom()) -> {ok, reference()} | {error, term()}.
bernoulli(_P, _Shape, _Dtype) ->
    erlang:nif_error({nif_not_loaded, ?MODULE}).

-spec categorical(reference(), [integer()], atom()) -> {ok, reference()} | {error, term()}.
categorical(_Logits, _Shape, _Dtype) ->
    erlang:nif_error({nif_not_loaded, ?MODULE}).

-spec multinomial(reference(), integer(), [integer()], atom()) -> {ok, reference()} | {error, term()}.
multinomial(_P, _N, _Shape, _Dtype) ->
    erlang:nif_error({nif_not_loaded, ?MODULE}).

-spec gamma(reference(), reference(), [integer()], atom()) -> {ok, reference()} | {error, term()}.
gamma(_Alpha, _Beta, _Shape, _Dtype) ->
    erlang:nif_error({nif_not_loaded, ?MODULE}).

-spec beta(reference(), reference(), [integer()], atom()) -> {ok, reference()} | {error, term()}.
beta(_Alpha, _Beta, _Shape, _Dtype) ->
    erlang:nif_error({nif_not_loaded, ?MODULE}).

-spec exponential(reference(), [integer()], atom()) -> {ok, reference()} | {error, term()}.
exponential(_Lambda, _Shape, _Dtype) ->
    erlang:nif_error({nif_not_loaded, ?MODULE}).

-spec poisson(reference(), [integer()], atom()) -> {ok, reference()} | {error, term()}.
poisson(_Lambda, _Shape, _Dtype) ->
    erlang:nif_error({nif_not_loaded, ?MODULE}).

-spec shuffle(reference(), integer()) -> {ok, reference()} | {error, term()}.
shuffle(_Array, _Axis) ->
    erlang:nif_error({nif_not_loaded, ?MODULE}).

-spec choice(reference(), integer(), boolean()) -> {ok, reference()} | {error, term()}.
choice(_Array, _Size, _Replace) ->
    erlang:nif_error({nif_not_loaded, ?MODULE}).

-spec permutation(integer()) -> {ok, reference()} | {error, term()}.
permutation(_N) ->
    erlang:nif_error({nif_not_loaded, ?MODULE}).

-spec permutation(reference(), integer()) -> {ok, reference()} | {error, term()}.
permutation(_Array, _Axis) ->
    erlang:nif_error({nif_not_loaded, ?MODULE}).