"""Lab055: единицы энергии, комплексный шум и раздельная проверка модели."""

from __future__ import annotations

import numpy as np
import pytest

from experiments import lab055_energy_audit as lab


def test_complex_power_and_chunk_accumulation():
    values = np.array([3 + 4j, -128 + 127j, 0j, -4 - 3j], dtype=np.complex64)
    whole, chunks = lab.Moments(), lab.Moments()
    whole.add(values)
    chunks.add(values[:1])
    chunks.add(values[1:])
    assert whole.result(20) == chunks.result(20)
    result = whole.result(20)
    assert result["power_code2"] == (25 + 128**2 + 127**2 + 25) / 4
    assert result["clipped_samples_fraction"] == .25
    assert result["clipped_components_fraction"] == .25
    assert result["symbol_energy_code2_seconds"] == result["power_code2"] / 20


def test_cs8_signed_layout_and_bounds(tmp_path):
    path = tmp_path / "input.cs8"
    path.write_bytes(np.array([-128, 127, 3, -4], dtype=np.int8).tobytes())
    source = lab.Cs8File(path)
    np.testing.assert_array_equal(source.read(0, 2), [-128 + 127j, 3 - 4j])
    with pytest.raises(ValueError):
        source.read(1, 2)
    with pytest.raises(ValueError):
        source.read(-1, 1)
    path = tmp_path / "odd.cs8"
    path.write_bytes(b"123")
    with pytest.raises(ValueError):
        lab.Cs8File(path)


@pytest.mark.parametrize("samples_per_symbol", [16, 400])
def test_complex_white_noise_convention_and_enbw(samples_per_symbol):
    taps, metrics = lab.filter_metrics(samples_per_symbol)
    variance = .7
    frequencies = np.fft.fftfreq(lab.SPECTRUM_SIZE, 1 / lab.SAMPLE_RATE)
    density = np.full(frequencies.size, variance / lab.SAMPLE_RATE)
    filtered = lab.matched_noise(frequencies, density, taps)
    assert filtered == pytest.approx(variance * sum(abs(taps)**2), rel=1e-10)
    assert filtered / (lab.SAMPLE_RATE * metrics["energy"]) == pytest.approx(variance / lab.SAMPLE_RATE)
    assert metrics["equivalent_noise_bandwidth_hz"] == pytest.approx(
        lab.SAMPLE_RATE * sum(abs(taps)**2) / abs(sum(taps))**2)


def test_psd_integrates_complex_power_without_factor_two():
    sample_rate = 65536
    values = 3 * np.exp(2j * np.pi * 4096 * np.arange(sample_rate) / sample_rate)
    frequencies, density = lab.spectrum(values, sample_rate)
    assert lab.band_power(frequencies, density, 4000, 4200) == pytest.approx(9, rel=1e-10)
    assert lab.band_power(frequencies, density, -4200, -4000) < 1e-20


def test_training_alignment_known_delay():
    random_generator = np.random.default_rng(5501)
    reference = random_generator.normal(size=1000) + 1j * random_generator.normal(size=1000)
    received = np.concatenate((np.zeros(12), reference, np.zeros(8))) + 4 + 2j
    assert lab.align_training(reference, received, 10) == 2


def test_channel_recovers_known_gain_frequency_and_heldout_noise():
    random_generator = np.random.default_rng(5502)
    count, sample_rate = 80000, 80000
    reference = random_generator.choice([-1., 1.], count) * np.exp(2j * np.pi * .1 * np.arange(count))
    gain, dc, frequency, variance = .3 * np.exp(.7j), 2 + 3j, .17, .02
    received = dc + gain * reference * np.exp(2j * np.pi * frequency * np.arange(count) / sample_rate)
    received += np.sqrt(variance / 2) * (random_generator.normal(size=count) + 1j * random_generator.normal(size=count))
    result = lab.fit_channel(reference, received, sample_rate)
    assert result["gain_abs"] == pytest.approx(abs(gain), rel=.02)
    assert result["cfo_hz"] == pytest.approx(frequency, abs=.01)
    assert result["heldout_signal_power"] == pytest.approx(.09, rel=.04)
    assert result["heldout_residual_power"] == pytest.approx(variance, rel=.04)
    assert result["heldout_actual_ac_power"] == pytest.approx(
        result["heldout_signal_power"] + result["heldout_residual_power"] + result["heldout_cross_power"], abs=1e-10)


def test_heldout_corruption_does_not_train_channel():
    random_generator = np.random.default_rng(5503)
    reference = random_generator.normal(size=8000) + 1j * random_generator.normal(size=8000)
    received = .2j * reference + 2
    clean = lab.fit_channel(reference, received, 8000)
    for tile in (1, 3, 5, 7):
        received[tile * 1000:(tile + 1) * 1000] += 5
    damaged = lab.fit_channel(reference, received, 8000)
    assert damaged["gain_abs"] == clean["gain_abs"]
    assert damaged["cfo_hz"] == clean["cfo_hz"]
    assert damaged["train_residual_power"] < 1e-20
    assert damaged["heldout_residual_power"] == pytest.approx(25)


def test_pipeline_timing_search_does_not_read_heldout_boundary():
    class MemorySource:
        def __init__(self, values):
            self.values = values

        def read(self, start, count):
            return self.values[start:start + count].copy()

    random_generator = np.random.default_rng(55055)
    count, samples_per_symbol, start = 2048, 16, 200
    reference = random_generator.normal(size=count) + 1j * random_generator.normal(size=count)
    received = np.zeros(count + 2 * start, dtype=complex)
    received[start:start + count] = reference
    clean = lab.analyze_rx(MemorySource(reference), MemorySource(received), [start], count, samples_per_symbol)
    first_heldout = start + 21 * samples_per_symbol + (count - 42 * samples_per_symbol) // 8
    received[first_heldout] += 10000
    damaged = lab.analyze_rx(MemorySource(reference), MemorySource(received), [start], count, samples_per_symbol)
    assert damaged["frames"][0]["alignment_adjustment"] == clean["frames"][0]["alignment_adjustment"] == 0
    assert damaged["frames"][0]["gain_abs"] == pytest.approx(1)
    assert damaged["frames"][0]["cfo_hz"] == pytest.approx(0)