mlx

-module(mlx_fft).

%% FFT and signal processing functions for MLX
-export([
    %% 1D FFT operations
    fft/1, fft/2,
    ifft/1, ifft/2,
    rfft/1, rfft/2,
    irfft/1, irfft/2,
    
    %% 2D FFT operations
    fft2/1, fft2/2,
    ifft2/1, ifft2/2,
    rfft2/1, rfft2/2,
    irfft2/1, irfft2/2,
    
    %% N-D FFT operations
    fftn/1, fftn/2, fftn/3,
    ifftn/1, ifftn/2, ifftn/3,
    rfftn/1, rfftn/2, rfftn/3,
    irfftn/1, irfftn/2, irfftn/3,
    
    %% FFT shift operations
    fftshift/1, fftshift/2,
    ifftshift/1, ifftshift/2,
    
    %% Frequency utilities
    fftfreq/2, fftfreq/3,
    rfftfreq/2, rfftfreq/3,
    
    %% Signal processing
    convolve/2, convolve/3,
    correlate/2, correlate/3,
    
    %% Window functions
    hann/1, hamming/1, blackman/1, bartlett/1,
    
    %% Spectral functions
    stft/2, stft/3, stft/4,
    istft/2, istft/3
]).

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

%% 1D FFT operations
-spec fft(array()) -> {ok, array()} | {error, term()}.
fft(Input) ->
    fft(Input, undefined).

-spec fft(array(), integer() | undefined) -> {ok, array()} | {error, term()}.
fft(Input, N) ->
    mlx_fft_nif:fft(Input, N).

-spec ifft(array()) -> {ok, array()} | {error, term()}.
ifft(Input) ->
    ifft(Input, undefined).

-spec ifft(array(), integer() | undefined) -> {ok, array()} | {error, term()}.
ifft(Input, N) ->
    mlx_fft_nif:ifft(Input, N).

-spec rfft(array()) -> {ok, array()} | {error, term()}.
rfft(Input) ->
    rfft(Input, undefined).

-spec rfft(array(), integer() | undefined) -> {ok, array()} | {error, term()}.
rfft(Input, N) ->
    mlx_fft_nif:rfft(Input, N).

-spec irfft(array()) -> {ok, array()} | {error, term()}.
irfft(Input) ->
    irfft(Input, undefined).

-spec irfft(array(), integer() | undefined) -> {ok, array()} | {error, term()}.
irfft(Input, N) ->
    mlx_fft_nif:irfft(Input, N).

%% 2D FFT operations
-spec fft2(array()) -> {ok, array()} | {error, term()}.
fft2(Input) ->
    fft2(Input, undefined).

-spec fft2(array(), shape() | undefined) -> {ok, array()} | {error, term()}.
fft2(Input, S) ->
    mlx_fft_nif:fft2(Input, S).

-spec ifft2(array()) -> {ok, array()} | {error, term()}.
ifft2(Input) ->
    ifft2(Input, undefined).

-spec ifft2(array(), shape() | undefined) -> {ok, array()} | {error, term()}.
ifft2(Input, S) ->
    mlx_fft_nif:ifft2(Input, S).

-spec rfft2(array()) -> {ok, array()} | {error, term()}.
rfft2(Input) ->
    rfft2(Input, undefined).

-spec rfft2(array(), shape() | undefined) -> {ok, array()} | {error, term()}.
rfft2(Input, S) ->
    mlx_fft_nif:rfft2(Input, S).

-spec irfft2(array()) -> {ok, array()} | {error, term()}.
irfft2(Input) ->
    irfft2(Input, undefined).

-spec irfft2(array(), shape() | undefined) -> {ok, array()} | {error, term()}.
irfft2(Input, S) ->
    mlx_fft_nif:irfft2(Input, S).

%% N-D FFT operations
-spec fftn(array()) -> {ok, array()} | {error, term()}.
fftn(Input) ->
    fftn(Input, undefined, undefined).

-spec fftn(array(), shape() | undefined) -> {ok, array()} | {error, term()}.
fftn(Input, S) ->
    fftn(Input, S, undefined).

-spec fftn(array(), shape() | undefined, [integer()] | undefined) -> {ok, array()} | {error, term()}.
fftn(Input, S, Axes) ->
    mlx_fft_nif:fftn(Input, S, Axes).

-spec ifftn(array()) -> {ok, array()} | {error, term()}.
ifftn(Input) ->
    ifftn(Input, undefined, undefined).

-spec ifftn(array(), shape() | undefined) -> {ok, array()} | {error, term()}.
ifftn(Input, S) ->
    ifftn(Input, S, undefined).

-spec ifftn(array(), shape() | undefined, [integer()] | undefined) -> {ok, array()} | {error, term()}.
ifftn(Input, S, Axes) ->
    mlx_fft_nif:ifftn(Input, S, Axes).

-spec rfftn(array()) -> {ok, array()} | {error, term()}.
rfftn(Input) ->
    rfftn(Input, undefined, undefined).

-spec rfftn(array(), shape() | undefined) -> {ok, array()} | {error, term()}.
rfftn(Input, S) ->
    rfftn(Input, S, undefined).

-spec rfftn(array(), shape() | undefined, [integer()] | undefined) -> {ok, array()} | {error, term()}.
rfftn(Input, S, Axes) ->
    mlx_fft_nif:rfftn(Input, S, Axes).

-spec irfftn(array()) -> {ok, array()} | {error, term()}.
irfftn(Input) ->
    irfftn(Input, undefined, undefined).

-spec irfftn(array(), shape() | undefined) -> {ok, array()} | {error, term()}.
irfftn(Input, S) ->
    irfftn(Input, S, undefined).

-spec irfftn(array(), shape() | undefined, [integer()] | undefined) -> {ok, array()} | {error, term()}.
irfftn(Input, S, Axes) ->
    mlx_fft_nif:irfftn(Input, S, Axes).

%% FFT shift operations
-spec fftshift(array()) -> {ok, array()} | {error, term()}.
fftshift(Input) ->
    fftshift(Input, undefined).

-spec fftshift(array(), [integer()] | undefined) -> {ok, array()} | {error, term()}.
fftshift(Input, Axes) ->
    mlx_fft_nif:fftshift(Input, Axes).

-spec ifftshift(array()) -> {ok, array()} | {error, term()}.
ifftshift(Input) ->
    ifftshift(Input, undefined).

-spec ifftshift(array(), [integer()] | undefined) -> {ok, array()} | {error, term()}.
ifftshift(Input, Axes) ->
    mlx_fft_nif:ifftshift(Input, Axes).

%% Frequency utilities
-spec fftfreq(integer(), number()) -> {ok, array()} | {error, term()}.
fftfreq(N, D) ->
    fftfreq(N, D, float32).

-spec fftfreq(integer(), number(), dtype()) -> {ok, array()} | {error, term()}.
fftfreq(N, D, Dtype) ->
    mlx_fft_nif:fftfreq(N, D, Dtype).

-spec rfftfreq(integer(), number()) -> {ok, array()} | {error, term()}.
rfftfreq(N, D) ->
    rfftfreq(N, D, float32).

-spec rfftfreq(integer(), number(), dtype()) -> {ok, array()} | {error, term()}.
rfftfreq(N, D, Dtype) ->
    mlx_fft_nif:rfftfreq(N, D, Dtype).

%% Signal processing
-spec convolve(array(), array()) -> {ok, array()} | {error, term()}.
convolve(Input, Kernel) ->
    convolve(Input, Kernel, full).

-spec convolve(array(), array(), atom()) -> {ok, array()} | {error, term()}.
convolve(Input, Kernel, Mode) ->
    mlx_fft_nif:convolve(Input, Kernel, Mode).

-spec correlate(array(), array()) -> {ok, array()} | {error, term()}.
correlate(Input, Kernel) ->
    correlate(Input, Kernel, full).

-spec correlate(array(), array(), atom()) -> {ok, array()} | {error, term()}.
correlate(Input, Kernel, Mode) ->
    mlx_fft_nif:correlate(Input, Kernel, Mode).

%% Window functions
-spec hann(integer()) -> {ok, array()} | {error, term()}.
hann(N) ->
    mlx_fft_nif:hann(N).

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

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

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

%% Spectral functions
-spec stft(array(), array()) -> {ok, array()} | {error, term()}.
stft(Input, Window) ->
    stft(Input, Window, 256, 128).

-spec stft(array(), array(), integer()) -> {ok, array()} | {error, term()}.
stft(Input, Window, NpFft) ->
    stft(Input, Window, NpFft, 128).

-spec stft(array(), array(), integer(), integer()) -> {ok, array()} | {error, term()}.
stft(Input, Window, NpFft, HopLength) ->
    mlx_fft_nif:stft(Input, Window, NpFft, HopLength).

-spec istft(array(), array()) -> {ok, array()} | {error, term()}.
istft(Stft, Window) ->
    istft(Stft, Window, 128).

-spec istft(array(), array(), integer()) -> {ok, array()} | {error, term()}.
istft(Stft, Window, HopLength) ->
    mlx_fft_nif:istft(Stft, Window, HopLength).