from __future__ import annotations

import inspect

import numpy as np
from scipy.signal import fftconvolve

from experiments import lab043_hackrf_bpsk as lab043
from experiments import lab044_hackrf_bpsk_packet as lab044
from experiments import lab048_receiver_performance as lab048
from protocol import bpsk_radio as radio


def test_preamble_reference_cache_preserves_exact_samples() -> None:
    lab044._preamble_reference_samples.cache_clear()
    cached = lab044._preamble_reference_samples()
    repeated = lab044._preamble_reference_samples()
    expected_waveform = lab044._scale_and_carrier(
        radio.bpsk_modulate(lab044.preamble_bits())
    )
    expected = lab044.signed_iq_bytes_to_complex64(
        lab044.lab043.to_cs8_bytes(expected_waveform)
    )

    assert repeated is cached
    assert cached.flags.writeable is False
    assert np.array_equal(cached, expected)


def test_matched_filter_taps_cache_preserves_exact_coefficients() -> None:
    lab043._matched_filter_taps.cache_clear()
    cached = lab043._matched_filter_taps()
    repeated = lab043._matched_filter_taps()
    expected = radio.root_raised_cosine_taps(
        0.35,
        lab043.SAMPLES_PER_SYMBOL,
        10,
    ).astype(np.float32)

    assert repeated is cached
    assert cached.flags.writeable is False
    assert np.array_equal(cached, expected)


def test_overlap_add_matched_filter_matches_previous_fft_result() -> None:
    random_generator = np.random.default_rng(48)
    samples = (
        random_generator.normal(size=50_000)
        + 1j * random_generator.normal(size=50_000)
    ).astype(np.complex64)
    start = 1_337
    length = 40_000
    segment = samples[start : start + length] - np.mean(samples)
    indexes = start + np.arange(len(segment))
    baseband = segment * np.exp(
        -2j
        * np.pi
        * lab043.CARRIER_OFFSET_HZ
        * indexes
        / lab043.SAMPLE_RATE_HZ
    )
    legacy = fftconvolve(
        baseband,
        radio.root_raised_cosine_taps(0.35, lab043.SAMPLES_PER_SYMBOL, 10),
        mode="full",
    )

    optimized, nominal = lab043._matched_waveform(samples, start, length)

    assert nominal == 16 * lab043.SAMPLES_PER_SYMBOL + 4_000
    assert optimized.dtype == np.complex64
    assert legacy.dtype == np.complex128
    maximum_difference = float(np.max(np.abs(optimized - legacy)))
    reference_peak = float(np.max(np.abs(legacy)))
    assert np.allclose(optimized, legacy, rtol=2e-6, atol=2e-6)
    assert maximum_difference / reference_peak < 5e-7


def test_carrier_downconversion_vector_cache_is_exact_and_bounded() -> None:
    lab043._carrier_downconversion_vector.cache_clear()
    start = 1_337
    length = 40_000
    cached = lab043._carrier_downconversion_vector(start, length)
    repeated = lab043._carrier_downconversion_vector(start, length)
    indexes = start + np.arange(length)
    expected = np.exp(
        -2j
        * np.pi
        * lab043.CARRIER_OFFSET_HZ
        * indexes
        / lab043.SAMPLE_RATE_HZ
    ).astype(np.complex64)

    assert repeated is cached
    assert cached.flags.writeable is False
    assert np.array_equal(cached, expected)
    for extra_length in range(length + 1, length + 5):
        lab043._carrier_downconversion_vector(start, extra_length)
    assert lab043._carrier_downconversion_vector.cache_info().currsize == 4


def test_specialized_preamble_timing_matches_previous_search() -> None:
    samples = lab044._preamble_reference_samples()
    matched, nominal = lab043._matched_waveform(samples, 0, samples.size)
    expected = radio.bpsk_modulate(lab044.preamble_bits())

    legacy, _ = lab043._timing_offsets(matched, nominal, expected)
    optimized = lab043._preamble_timing_offset(matched, nominal, expected)

    assert optimized == legacy


def test_profiled_receiver_api_has_no_tx_or_expected_frame_oracle() -> None:
    parameters = inspect.signature(lab048.run_profiled_receiver).parameters

    assert tuple(parameters) == ("capture_path", "original_result_path")
    assert "expected_frame_count" not in parameters
    assert "tx_image_path" not in parameters


def test_lab048_functional_checks() -> None:
    lab048.run_functional_tests()