import numpy as np

from rfg_adc_plotter.processing.fourier import (
    build_frequency_axis_hz,
    compute_ifft_profile_from_sweep,
    normalize_sweep_for_phase,
    perform_ifft_depth_response,
    reconstruct_complex_spectrum_from_real_trace,
    unwrap_arccos_phase_continuous,
)


def test_normalize_sweep_for_phase_max_abs_and_finite():
    sweep = np.array([np.nan, -10.0, 5.0, 20.0, -40.0, np.inf, -np.inf], dtype=np.float32)
    x = normalize_sweep_for_phase(sweep)
    assert x.dtype == np.float64
    assert np.all(np.isfinite(x))
    assert np.max(np.abs(x)) <= 1.0 + 1e-12


def test_arccos_unwrap_continuous_recovers_complex_phase_without_large_jumps():
    phi_true = np.linspace(0.0, 4.0 * np.pi, 1000, dtype=np.float64)
    x = np.cos(phi_true)

    phi_rec = unwrap_arccos_phase_continuous(x)
    assert phi_rec.shape == phi_true.shape
    assert np.max(np.abs(np.diff(phi_rec))) < 0.2

    z_true = np.exp(1j * phi_true)
    z_rec = np.exp(1j * phi_rec)
    assert np.allclose(z_rec, z_true, atol=2e-2, rtol=0.0)


def test_reconstruct_complex_spectrum_from_real_trace_output_complex128():
    sweep = np.linspace(-1.0, 1.0, 64, dtype=np.float32)
    z = reconstruct_complex_spectrum_from_real_trace(sweep)
    assert z.dtype == np.complex128
    assert z.shape == sweep.shape
    assert np.all(np.isfinite(np.real(z)))
    assert np.all(np.isfinite(np.imag(z)))


def test_perform_ifft_depth_response_basic_abs():
    n = 128
    freqs = build_frequency_axis_hz(n)
    s = np.exp(1j * np.linspace(0.0, 2.0 * np.pi, n, dtype=np.float64))

    depth_m, y = perform_ifft_depth_response(s, freqs, axis="abs")

    assert depth_m.dtype == np.float32
    assert y.dtype == np.float32
    assert depth_m.ndim == 1 and y.ndim == 1
    assert depth_m.size == y.size
    assert depth_m.size >= n
    assert np.all(np.diff(depth_m) >= 0.0)
    assert np.all(y >= 0.0)


def test_perform_ifft_depth_response_bad_grid_returns_fallback_not_exception():
    s = np.ones(16, dtype=np.complex128)
    freqs_desc = np.linspace(10.0, 1.0, 16, dtype=np.float64)
    depth_m, y = perform_ifft_depth_response(s, freqs_desc, axis="abs")
    assert depth_m.size == y.size
    assert depth_m.size == s.size
    assert np.all(np.isfinite(depth_m))


def test_compute_ifft_profile_from_sweep_returns_depth_and_linear_abs():
    sweep = np.linspace(-5.0, 7.0, 257, dtype=np.float32)
    depth_m, y = compute_ifft_profile_from_sweep(sweep)
    assert depth_m.dtype == np.float32
    assert y.dtype == np.float32
    assert depth_m.size == y.size
    assert depth_m.size > 0
    assert np.all(np.diff(depth_m) >= 0.0)