import numpy as np

from galileo_rx.fixed_point.fixed_fft import FixedRadix2Fft
from galileo_rx.fixed_point.formats import FixedFormat


def test_fixed_fft_impulse_and_scaling() -> None:
    values = np.zeros(16, dtype=np.complex128)
    values[0] = 0.5
    model = FixedRadix2Fft(FixedFormat(20, 17))
    result = model.forward(values)
    np.testing.assert_allclose(result.dequantize(), np.fft.fft(values) / values.size, atol=4 * model.data_format.lsb)
    assert result.overflow_count == 0


def test_fixed_fft_random_vector_and_inverse() -> None:
    rng = np.random.default_rng(91)
    values = 0.1 * (rng.standard_normal(32) + 1j * rng.standard_normal(32))
    model = FixedRadix2Fft(FixedFormat(24, 19))
    spectrum = model.forward(values).dequantize()
    np.testing.assert_allclose(spectrum, np.fft.fft(values) / values.size, atol=2e-5)
    restored = model.inverse(spectrum).dequantize()
    # IFFT of the already 1/N-scaled forward spectrum is x/N.
    np.testing.assert_allclose(restored, values / values.size, atol=2e-5)


def test_fixed_fft_rejects_non_power_of_two() -> None:
    import pytest

    with pytest.raises(ValueError, match="power of two"):
        FixedRadix2Fft().forward(np.ones(12))