"""Быстрые тесты диагностических примитивов Lab058."""

from __future__ import annotations

import numpy as np
import pytest

from experiments import lab058_abab_symbol_quality as lab058


def test_training_and_validation_tiles_are_disjoint_and_complete():
    for count in (4096, 102400):
        train = lab058.training_mask(count)
        assert train.dtype == bool and train.size == count
        assert np.all((lab058.tile_ids(count)[train] % 2) == 0)
        assert np.all((lab058.tile_ids(count)[~train] % 2) == 1)


def test_guarded_tiles_are_disjoint_and_exclude_boundaries():
    for count in (4096, 102400):
        train, validation = lab058.guarded_tile_masks(count)
        assert not np.any(train & validation)
        assert train.sum() == validation.sum() == count // 2 - 16 * 32


@pytest.mark.parametrize("symbol_rate", [20_000, 500_000])
def test_cfo_gain_fit_recovers_known_symbol_model(symbol_rate):
    expected = lab058.expected_symbols(symbol_rate)[:4096]
    indexes = np.arange(len(expected))
    cfo = 1.75 * symbol_rate / 20_000
    h = .7 - .25j
    samples = h * expected * np.exp(2j * np.pi * cfo * (indexes + .5) / symbol_rate)
    found_cfo, found_h, coherence = lab058.fit_cfo_gain(samples, expected, indexes, symbol_rate)
    assert found_cfo == pytest.approx(cfo, abs=1e-8)
    assert found_h == pytest.approx(h, abs=1e-10)
    assert coherence == pytest.approx(1.0)


def test_cfo_gain_subset_keeps_original_tile_time_axis():
    symbol_rate = 500_000
    full = lab058.expected_symbols(symbol_rate)
    indexes = np.flatnonzero(lab058.training_mask(len(full)))[::31]
    cfo, h = 3.25, .4 + .2j
    samples = h * full[indexes] * np.exp(2j * np.pi * cfo * (indexes + .5) / symbol_rate)
    found, gain, _ = lab058.fit_cfo_gain(samples, full[indexes], indexes, symbol_rate, len(full))
    assert found == pytest.approx(cfo, abs=1e-8)
    assert gain == pytest.approx(h, abs=1e-9)


def test_oracle_estimate_does_not_read_guarded_validation_symbols():
    symbol_rate = 500_000
    expected = lab058.expected_symbols(symbol_rate)[:4096]
    samples_per_symbol = lab058.SAMPLE_RATE // symbol_rate
    values = lab058.shaping.shape_frame_like_lab042(expected, samples_per_symbol)
    values *= np.exp(2j * np.pi * lab058.CARRIER_HZ * np.arange(len(values)) / lab058.SAMPLE_RATE)
    changed = values.copy()
    validation = np.flatnonzero(lab058.guarded_tile_masks(len(expected))[1])
    centers = 21 * samples_per_symbol + validation * samples_per_symbol
    for center in centers:
        changed[center:center + samples_per_symbol] += 3 - 2j
    frontend = lab058.lab057.Frontend(symbol_rate)
    first_matched = lab058.oracle_matched_filter(values, frontend.taps, symbol_rate, len(expected))
    second_matched = lab058.oracle_matched_filter(changed, frontend.taps, symbol_rate, len(expected))
    first = lab058.oracle_estimate(first_matched, frontend.nominal, expected, symbol_rate)
    second = lab058.oracle_estimate(second_matched, frontend.nominal, expected, symbol_rate)
    assert first["timing_fraction"] == second["timing_fraction"]
    assert first["cfo_hz"] == pytest.approx(second["cfo_hz"], abs=1e-5)
    assert first["h"] == pytest.approx(second["h"], abs=1e-7)
    assert first["train_coherence"] == pytest.approx(second["train_coherence"], abs=1e-8)


def test_evm_and_margin_formula_on_exact_symbols():
    expected = lab058.expected_symbols(20_000)
    indexes = np.arange(len(expected))
    h = .4 + .3j
    estimate = {"symbols": h * expected, "h": h, "cfo_hz": 0.0}
    result = lab058.evaluate(estimate, expected, 20_000, np.ones(len(expected), bool))
    assert result["errors"] == 0 and result["ber"] == 0
    assert result["evm_rms"] < 1e-14
    assert result["margin_median"] == pytest.approx(1.0)


def test_isi_fit_localizes_single_main_tap():
    expected = lab058.expected_symbols(20_000)
    h = .2 - .1j
    estimate = {"symbols": h * expected, "h": h, "cfo_hz": 0.0}
    result, taps = lab058.isi_metrics(estimate, expected, 20_000)
    assert abs(taps[5] - h) < 1e-12
    assert result["adjacent_taps_power"] < result["main_tap_power"] * 1e-20


def test_averaged_transfer_estimate_is_flat_for_flat_synthetic_channel():
    random_generator = np.random.default_rng(5802)
    count = 131072
    tx = (random_generator.normal(size=count) + 1j * random_generator.normal(size=count)).astype(np.complex64)
    gain = .3 - .2j
    result, _ = lab058.channel_metrics(tx, gain * tx, 500_000, 0.0)
    assert result["stable"]
    assert result["median_coherence"] > .999
    assert result["amplitude_ripple_p95_p05_db"] < .01