viva_math

-module(viva_math@fft).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_math/fft.gleam").
-export([fft_size/1, fft/1, ifft/1, pad_to_power_of_two/1]).

-if(?OTP_RELEASE >= 27).
-define(MODULEDOC(Str), -moduledoc(Str)).
-define(DOC(Str), -doc(Str)).
-else.
-define(MODULEDOC(Str), -compile([])).
-define(DOC(Str), -compile([])).
-endif.

?MODULEDOC(
    " Fast Fourier Transform over `viva_math/complex`.\n"
    "\n"
    " Implements the recursive Cooley-Tukey radix-2 FFT. Inputs to `fft` and\n"
    " `ifft` must have power-of-two length; use `pad_to_power_of_two` when a\n"
    " signal needs zero-padding first.\n"
    "\n"
    " ## Examples\n"
    "\n"
    " ```gleam\n"
    " import viva_math/complex\n"
    " import viva_math/fft\n"
    "\n"
    " fft.fft([\n"
    "   complex.real(1.0),\n"
    "   complex.zero(),\n"
    "   complex.zero(),\n"
    "   complex.zero(),\n"
    " ])\n"
    " // -> [1+0i, 1+0i, 1+0i, 1+0i]\n"
    " ```\n"
).

-file("src/viva_math/fft.gleam", 106).
-spec combine(
    list(viva_math@complex:complex()),
    list(viva_math@complex:complex()),
    integer(),
    integer(),
    list(viva_math@complex:complex()),
    list(viva_math@complex:complex())
) -> list(viva_math@complex:complex()).
combine(Evens, Odds, K, N, Low, High) ->
    case {Evens, Odds} of
        {[], []} ->
            lists:append(lists:reverse(Low), lists:reverse(High));

        {[E | Even_rest], [O | Odd_rest]} ->
            Angle = +0.0 - (case erlang:float(N) of
                +0.0 -> +0.0;
                -0.0 -> -0.0;
                Gleam@denominator -> 6.283185307179586 * erlang:float(K) / Gleam@denominator
            end),
            Twiddle = viva_math@complex:from_polar(1.0, Angle),
            T = viva_math@complex:mul(Twiddle, O),
            combine(
                Even_rest,
                Odd_rest,
                K + 1,
                N,
                [viva_math@complex:add(E, T) | Low],
                [viva_math@complex:sub(E, T) | High]
            );

        {_, _} ->
            lists:append(lists:reverse(Low), lists:reverse(High))
    end.

-file("src/viva_math/fft.gleam", 94).
-spec split_even_odd(
    list(viva_math@complex:complex()),
    list(viva_math@complex:complex()),
    list(viva_math@complex:complex())
) -> {list(viva_math@complex:complex()), list(viva_math@complex:complex())}.
split_even_odd(Signal, Evens, Odds) ->
    case Signal of
        [] ->
            {lists:reverse(Evens), lists:reverse(Odds)};

        [Even] ->
            {lists:reverse([Even | Evens]), lists:reverse(Odds)};

        [Even@1, Odd | Rest] ->
            split_even_odd(Rest, [Even@1 | Evens], [Odd | Odds])
    end.

-file("src/viva_math/fft.gleam", 76).
-spec fft_unchecked(list(viva_math@complex:complex())) -> list(viva_math@complex:complex()).
fft_unchecked(Signal) ->
    case Signal of
        [] ->
            [];

        [_] ->
            Signal;

        _ ->
            {Even, Odd} = split_even_odd(Signal, [], []),
            combine(
                fft_unchecked(Even),
                fft_unchecked(Odd),
                0,
                erlang:length(Signal),
                [],
                []
            )
    end.

-file("src/viva_math/fft.gleam", 130).
-spec is_power_of_two(integer()) -> boolean().
is_power_of_two(N) ->
    case N of
        N@1 when N@1 =< 0 ->
            false;

        1 ->
            true;

        N@2 when (N@2 rem 2) =/= 0 ->
            false;

        _ ->
            is_power_of_two(N div 2)
    end.

-file("src/viva_math/fft.gleam", 54).
?DOC(
    " Return the signal length if it is a power of two.\n"
    "\n"
    " Panics for empty or non-power-of-two inputs because the public API returns\n"
    " `Int` directly.\n"
).
-spec fft_size(list(viva_math@complex:complex())) -> integer().
fft_size(Signal) ->
    N = erlang:length(Signal),
    case is_power_of_two(N) of
        true ->
            N;

        false ->
            erlang:error(#{gleam_error => panic,
                    message => <<"FFT input length must be a non-zero power of two"/utf8>>,
                    file => <<?FILEPATH/utf8>>,
                    module => <<"viva_math/fft"/utf8>>,
                    function => <<"fft_size"/utf8>>,
                    line => 58})
    end.

-file("src/viva_math/fft.gleam", 30).
?DOC(
    " Forward FFT.\n"
    "\n"
    " The signal length must be a power of two. Invalid lengths panic via\n"
    " `fft_size`, matching the total return type requested for this module.\n"
).
-spec fft(list(viva_math@complex:complex())) -> list(viva_math@complex:complex()).
fft(Signal) ->
    _ = fft_size(Signal),
    fft_unchecked(Signal).

-file("src/viva_math/fft.gleam", 40).
?DOC(
    " Inverse FFT using `IFFT(x) = conj(FFT(conj(x))) / N`.\n"
    "\n"
    " The spectrum length must be a power of two. The final result is normalized\n"
    " by `N`, so `ifft(fft(signal))` recovers the original signal up to floating\n"
    " point round-off.\n"
).
-spec ifft(list(viva_math@complex:complex())) -> list(viva_math@complex:complex()).
ifft(Spectrum) ->
    N = fft_size(Spectrum),
    Scale = case erlang:float(N) of
        +0.0 -> +0.0;
        -0.0 -> -0.0;
        Gleam@denominator -> 1.0 / Gleam@denominator
    end,
    _pipe = Spectrum,
    _pipe@1 = gleam@list:map(_pipe, fun viva_math@complex:conjugate/1),
    _pipe@2 = fft_unchecked(_pipe@1),
    gleam@list:map(
        _pipe@2,
        fun(Z) ->
            viva_math@complex:scale(viva_math@complex:conjugate(Z), Scale)
        end
    ).

-file("src/viva_math/fft.gleam", 143).
-spec next_power_of_two_loop(integer(), integer()) -> integer().
next_power_of_two_loop(N, Acc) ->
    case Acc >= N of
        true ->
            Acc;

        false ->
            next_power_of_two_loop(N, Acc * 2)
    end.

-file("src/viva_math/fft.gleam", 139).
-spec next_power_of_two(integer()) -> integer().
next_power_of_two(N) ->
    next_power_of_two_loop(N, 1).

-file("src/viva_math/fft.gleam", 65).
?DOC(
    " Pad a signal with complex zeros up to the next power of two.\n"
    "\n"
    " Already valid lengths are returned unchanged; an empty signal remains empty.\n"
).
-spec pad_to_power_of_two(list(viva_math@complex:complex())) -> list(viva_math@complex:complex()).
pad_to_power_of_two(Signal) ->
    N = erlang:length(Signal),
    case N of
        0 ->
            [];

        _ ->
            Target = next_power_of_two(N),
            lists:append(
                Signal,
                gleam@list:repeat(viva_math@complex:zero(), Target - N)
            )
    end.