from __future__ import annotations

import hashlib
import math
from pathlib import Path

import numpy as np
import pytest
from scipy import signal

from tools.rf001_staircase_search import (
    DEFAULT_CRITERIA,
    DEFAULT_LAYOUT,
    EXPECTED_REFERENCE_SHA256,
    ReferenceLayout,
    SearchCriteria,
    analyze,
    design_antialias_filter,
    locate_reference_streaming,
    load_reference,
    read_filtered_decimated_window,
    verify_candidate,
)


LAYOUT = ReferenceLayout(
    sample_rate_hz=8_000,
    symbol_rate=200,
    carrier_hz=500.0,
    leading_samples=800,
    segment_samples=800,
    gap_samples=400,
    trailing_samples=800,
)
CRITERIA = SearchCriteria(
    decimation=1,
    input_block_samples=50_000,
    cfo_min_hz=-3.0,
    cfo_max_hz=3.0,
    cfo_step_hz=0.5,
    histogram_stride=8,
)
PROJECT_ROOT = Path(__file__).resolve().parents[1]
EXACT_REFERENCE_PATH = (
    PROJECT_ROOT / ".local" / "rf001_probe" / "lab065_reference_staircase.cs8"
)


def _reference() -> np.ndarray:
    random_generator = np.random.default_rng(20260925)
    symbols = random_generator.choice((-1.0, 1.0), LAYOUT.segment_samples // 40)
    base = np.repeat(symbols, 40).astype(np.complex64)
    time_axis = np.arange(LAYOUT.segment_samples) / LAYOUT.sample_rate_hz
    base *= np.exp(2j * np.pi * LAYOUT.carrier_hz * time_axis)
    base *= 60.0 / np.sqrt(np.mean(np.abs(base) ** 2))
    values = np.zeros(LAYOUT.total_samples, dtype=np.complex64)
    for level, (start, stop) in zip(LAYOUT.levels_db, LAYOUT.active_windows()):
        values[start:stop] = base * (10 ** (level / 20))
    return values


def _write_cs8(path: Path, values: np.ndarray) -> str:
    interleaved = np.empty(2 * values.size, dtype=np.int8)
    interleaved[0::2] = np.clip(np.rint(values.real), -128, 127).astype(np.int8)
    interleaved[1::2] = np.clip(np.rint(values.imag), -128, 127).astype(np.int8)
    path.write_bytes(interleaved.tobytes())
    return hashlib.sha256(interleaved.tobytes()).hexdigest()


def _capture(reference: np.ndarray, *, gain: float, cfo_hz: float, start: int) -> np.ndarray:
    random_generator = np.random.default_rng(417)
    values = (
        random_generator.normal(0.0, 1.5, 1_000_000)
        + 1j * random_generator.normal(0.0, 1.5, 1_000_000)
    ).astype(np.complex64)
    time_axis = np.arange(reference.size) / LAYOUT.sample_rate_hz
    values[start:start + reference.size] += (
        gain
        * reference
        * np.exp(1j * (0.73 + 2 * np.pi * cfo_hz * time_axis))
    )
    return values


@pytest.mark.parametrize(
    ("gain", "cfo_hz", "start"),
    ((0.25, -2.5, 49_777), (0.75, 0.0, 420_201), (1.5, 2.5, 749_999)),
)
def test_unknown_position_gain_cfo_and_block_boundary(
    tmp_path: Path,
    gain: float,
    cfo_hz: float,
    start: int,
) -> None:
    reference_path = tmp_path / "reference.cs8"
    capture_path = tmp_path / "capture.cs8"
    reference_hash = _write_cs8(reference_path, _reference())
    _write_cs8(capture_path, _capture(_reference(), gain=gain, cfo_hz=cfo_hz, start=start))
    reference, _ = load_reference(
        reference_path,
        LAYOUT,
        CRITERIA,
        expected_sha256=reference_hash,
    )
    search = locate_reference_streaming(capture_path, reference, LAYOUT, CRITERIA)
    candidate = search["candidates"][0]
    verification = verify_candidate(capture_path, reference, candidate, LAYOUT, CRITERIA)
    assert abs(int(candidate["start_sample"]) - start) <= 1
    assert abs(float(candidate["cfo_hz"]) - cfo_hz) <= CRITERIA.cfo_step_hz
    assert float(candidate["score"]) >= CRITERIA.normalized_correlation_minimum
    assert verification["qualified"] is True


def test_negative_control_without_staircase_is_not_detected(tmp_path: Path) -> None:
    reference_path = tmp_path / "reference.cs8"
    capture_path = tmp_path / "noise.cs8"
    report_path = tmp_path / "report.json"
    reference_hash = _write_cs8(reference_path, _reference())
    random_generator = np.random.default_rng(818)
    noise = (
        random_generator.normal(0.0, 3.0, 1_000_000)
        + 1j * random_generator.normal(0.0, 3.0, 1_000_000)
    ).astype(np.complex64)
    capture_hash = _write_cs8(capture_path, noise)
    result = analyze(
        reference_path,
        capture_path,
        report_path,
        expected_capture_sha256=capture_hash,
        layout=LAYOUT,
        criteria=CRITERIA,
        expected_reference_sha256=reference_hash,
    )
    assert result["verdict"] == "NOT_DETECTED"
    assert report_path.is_file()
    assert result["hardware_accessed"] is False
    assert result["rf_tx"] == 0


def test_report_is_not_overwritten(tmp_path: Path) -> None:
    reference_path = tmp_path / "reference.cs8"
    capture_path = tmp_path / "capture.cs8"
    report_path = tmp_path / "report.json"
    reference_hash = _write_cs8(reference_path, _reference())
    capture_hash = _write_cs8(
        capture_path,
        _capture(_reference(), gain=1.0, cfo_hz=0.0, start=4_000),
    )
    result = analyze(
        reference_path,
        capture_path,
        report_path,
        expected_capture_sha256=capture_hash,
        layout=LAYOUT,
        criteria=CRITERIA,
        expected_reference_sha256=reference_hash,
    )
    assert result["verdict"] == "DETECTED"
    with pytest.raises(FileExistsError):
        analyze(
            reference_path,
            capture_path,
            report_path,
            expected_capture_sha256=capture_hash,
            layout=LAYOUT,
            criteria=CRITERIA,
            expected_reference_sha256=reference_hash,
        )


def _require_exact_reference() -> Path:
    if not EXACT_REFERENCE_PATH.is_file():
        pytest.skip("exact Lab065 reference is not available in this checkout")
    assert EXACT_REFERENCE_PATH.stat().st_size == 12_000_000
    assert hashlib.sha256(EXACT_REFERENCE_PATH.read_bytes()).hexdigest() == (
        EXPECTED_REFERENCE_SHA256
    )
    return EXACT_REFERENCE_PATH


def _read_exact_reference_values(path: Path) -> np.ndarray:
    raw = np.memmap(path, dtype=np.int8, mode="r")
    return (
        np.asarray(raw[0::2], dtype=np.float32)
        + 1j * np.asarray(raw[1::2], dtype=np.float32)
    ).astype(np.complex64)


def _write_noise_cs8(path: Path, samples: int, seed: int) -> None:
    random_generator = np.random.default_rng(seed)
    raw = np.memmap(path, dtype=np.int8, mode="w+", shape=(2 * samples,))
    block = 500_000
    for start in range(0, samples, block):
        stop = min(samples, start + block)
        count = stop - start
        raw[2 * start:2 * stop:2] = np.clip(
            np.rint(random_generator.normal(0.0, 0.75, count)),
            -128,
            127,
        ).astype(np.int8)
        raw[2 * start + 1:2 * stop + 1:2] = np.clip(
            np.rint(random_generator.normal(0.0, 0.75, count)),
            -128,
            127,
        ).astype(np.int8)
    raw.flush()
    del raw


def _inject_exact_reference(
    capture_path: Path,
    reference: np.ndarray,
    *,
    start: int,
    gain: float,
    phase: float,
    cfo_hz: float,
) -> None:
    raw = np.memmap(capture_path, dtype=np.int8, mode="r+")
    existing = (
        np.asarray(raw[2 * start:2 * (start + reference.size):2], dtype=np.float32)
        + 1j
        * np.asarray(
            raw[2 * start + 1:2 * (start + reference.size) + 1:2],
            dtype=np.float32,
        )
    )
    time_axis = np.arange(reference.size, dtype=np.float64) / DEFAULT_LAYOUT.sample_rate_hz
    injected = existing + gain * reference * np.exp(
        1j * (phase + 2 * np.pi * cfo_hz * time_axis)
    )
    raw[2 * start:2 * (start + reference.size):2] = np.clip(
        np.rint(injected.real), -128, 127
    ).astype(np.int8)
    raw[2 * start + 1:2 * (start + reference.size) + 1:2] = np.clip(
        np.rint(injected.imag), -128, 127
    ).astype(np.int8)
    raw.flush()
    del raw


def test_antialias_filter_and_block_overlap_are_fixed() -> None:
    reference_path = _require_exact_reference()
    taps = design_antialias_filter(DEFAULT_LAYOUT, DEFAULT_CRITERIA)
    assert taps.size == 895
    assert (taps.size - 1) // 2 == 447
    frequencies, response = signal.freqz(
        taps,
        worN=262_144,
        fs=DEFAULT_LAYOUT.sample_rate_hz,
    )
    assert np.max(np.abs(response[frequencies >= 120_000])) <= 10 ** (-79 / 20)

    total = math.ceil(DEFAULT_LAYOUT.total_samples / DEFAULT_CRITERIA.decimation)
    split = 73_219
    whole = read_filtered_decimated_window(
        reference_path,
        0,
        total,
        DEFAULT_LAYOUT,
        DEFAULT_CRITERIA,
        taps,
    )
    blocked = np.concatenate(
        (
            read_filtered_decimated_window(
                reference_path,
                0,
                split,
                DEFAULT_LAYOUT,
                DEFAULT_CRITERIA,
                taps,
            ),
            read_filtered_decimated_window(
                reference_path,
                split,
                total,
                DEFAULT_LAYOUT,
                DEFAULT_CRITERIA,
                taps,
            ),
        )
    )
    np.testing.assert_allclose(blocked, whole, rtol=0, atol=2e-5)


def test_exact_reference_unknown_positions_gain_phase_noise_and_cfo(
    tmp_path: Path,
) -> None:
    reference_path = _require_exact_reference()
    reference_values = _read_exact_reference_values(reference_path)
    cases = (
        (123_457, 0.45, 0.73, -2.5),
        (7_999_987, 0.70, -1.17, 0.0),
        (15_999_983, 1.00, 2.21, 2.5),
    )
    capture_path = tmp_path / "exact_reference_cases.cs8"
    _write_noise_cs8(capture_path, 22_500_000, seed=2026092801)
    for start, gain, phase, cfo_hz in cases:
        _inject_exact_reference(
            capture_path,
            reference_values,
            start=start,
            gain=gain,
            phase=phase,
            cfo_hz=cfo_hz,
        )
    reference, _ = load_reference(
        reference_path,
        DEFAULT_LAYOUT,
        DEFAULT_CRITERIA,
        expected_sha256=EXPECTED_REFERENCE_SHA256,
    )
    search = locate_reference_streaming(
        capture_path,
        reference,
        DEFAULT_LAYOUT,
        DEFAULT_CRITERIA,
    )
    assert search["coverage_complete"] is True
    assert search["source_samples_processed"] == 22_500_000
    candidates = list(search["candidates"])
    for expected_start, _gain, _phase, expected_cfo in cases:
        candidate = min(
            candidates,
            key=lambda row: abs(int(row["start_sample"]) - expected_start),
        )
        assert abs(int(candidate["start_sample"]) - expected_start) <= 32
        assert abs(float(candidate["cfo_hz"]) - expected_cfo) <= 0.25
        assert float(candidate["score"]) >= DEFAULT_CRITERIA.normalized_correlation_minimum
        verification = verify_candidate(
            capture_path,
            reference,
            candidate,
            DEFAULT_LAYOUT,
            DEFAULT_CRITERIA,
        )
        assert verification["qualified"] is True


def test_exact_reference_negative_control(tmp_path: Path) -> None:
    reference_path = _require_exact_reference()
    capture_path = tmp_path / "exact_reference_negative.cs8"
    report_path = tmp_path / "negative_report.json"
    _write_noise_cs8(capture_path, 8_000_000, seed=2026092802)
    capture_hash = hashlib.sha256(capture_path.read_bytes()).hexdigest()
    result = analyze(
        reference_path,
        capture_path,
        report_path,
        expected_capture_sha256=capture_hash,
    )
    assert result["verdict"] == "NOT_DETECTED"
    assert result["search"]["coverage_complete"] is True