import numpy as np

from experiments import lab043_bpsk_asymmetry_diagnostics as diagnostics


def test_known_sequence_phase_fit_removes_known_linear_phase() -> None:
    bits = np.tile(np.array([0, 0, 1, 1, 0, 1, 0, 1], dtype=np.uint8), 128)
    expected = 2.0 * bits.astype(np.float64) - 1.0
    indexes = np.arange(len(expected), dtype=np.float64)
    frequency_hz = 0.75
    received = expected * np.exp(
        1j * (0.37 + 2.0 * np.pi * frequency_hz * indexes / 20_000.0)
    )

    corrected, estimated_hz, fit_rmse = diagnostics._known_sequence_phase_fit(
        received, expected, 20_000.0
    )
    normalized, _ = diagnostics._normalize_symbols(corrected, expected)

    assert abs(estimated_hz - frequency_hz) < 1e-9
    assert fit_rmse < 1e-12
    assert np.max(np.abs(normalized - expected)) < 1e-12


def test_triplet_rows_count_and_classify_only_actual_decision_crossings() -> None:
    bits = np.tile(np.array([0, 0, 0, 1, 1, 1, 0, 1], dtype=np.uint8), 20)
    values = 2.0 * bits.astype(np.float64) - 1.0
    target = np.flatnonzero(
        (bits[:-2] == 0) & (bits[1:-1] == 0) & (bits[2:] == 0)
    )[0] + 1
    values[target] = 0.1

    rows = diagnostics.triplet_rows("synthetic", bits, values.astype(np.complex128))
    by_triplet = {row["triplet"]: row for row in rows}

    assert sum(row["count"] for row in rows) == len(bits) - 2
    assert by_triplet["000"]["bit_errors"] == 1
    assert sum(row["bit_errors"] for row in rows) == 1
    assert by_triplet["000"]["decision_cross_fraction"] == by_triplet["000"]["ber"]


def test_rrc_audit_has_symmetric_unit_energy_pair() -> None:
    impulse_rows, audit = diagnostics.rrc_rows()

    assert audit["tx_tap_count"] == 1201
    assert audit["rx_tap_count"] == 1201
    assert audit["symmetry_max_abs_error"] == 0.0
    assert abs(audit["tap_energy"] - 1.0) < 1e-12
    by_offset = {row["symbol_offset"]: row["h_normalized"] for row in impulse_rows}
    assert by_offset[0] == 1.0
    for offset in range(1, 6):
        assert abs(by_offset[offset] - by_offset[-offset]) < 1e-15


def test_fractional_delay_preserves_order_and_length() -> None:
    source = np.array([0.0, 1.0, 2.0, 3.0], dtype=np.complex128)

    shifted = diagnostics._fractional_delay(source, 0.5)

    assert len(shifted) == len(source)
    assert np.allclose(shifted, [0.0, 0.5, 1.5, 2.5])


def test_triplet_metrics_keep_complex_dc_visible() -> None:
    bits = np.tile(np.array([0, 0, 1, 1, 0, 1, 0, 1], dtype=np.uint8), 20)
    expected = 2.0 * bits.astype(np.float64) - 1.0
    offset = 0.2 - 0.1j

    rows = diagnostics.triplet_rows(
        "affine", bits, expected.astype(np.complex128) + offset
    )

    negative = [row for row in rows if row["triplet"][1] == "0" and row["count"]]
    positive = [row for row in rows if row["triplet"][1] == "1" and row["count"]]
    present = [row for row in rows if row["count"]]
    assert all(abs(row["mean_i"] + 0.8) < 1e-12 for row in negative)
    assert all(abs(row["mean_i"] - 1.2) < 1e-12 for row in positive)
    assert all(abs(row["mean_q"] + 0.1) < 1e-12 for row in present)