scanengine-3/tests/test_saw_check.py

"""Middle-row SAW quality check: plan reduction, the v11 file, and the read-out.

The acquisition half runs on the same fake rig as the scan tests; the
analysis half runs on a synthetic v11 file whose CH1 is a pure sine at a
known FFT bin, so the frequency a trace reports is a number the test knows
in advance rather than one it copies from the implementation.
"""
import math

import numpy as np
import pytest

from core.rotation import RotationAxis, RotationSettings
from core.saw_check import (
    SPREAD_GOOD_PCT, alignment_summary, frequency_traces, middle_row_index,
    middle_row_plan,
)
from core.scan_engine import ScanCallbacks, ScanEngine
from core.scan_geometry import ScanGeometryError, build_plan
from core.sras_format import (
    SCAN_CHANNELS, VERSION, VERSION_SAW_CHECK, SrasFile, create_scan_file,
    write_background_block,
)
from fakes import FakeScope, FakeStage, FakeT3R, Trace

SAMPLE_RATE = 6.25e9
SPF = 256
LASER_FREQ_HZ = 20000.0
VELOCITY_MM_S = 100.0
PREAMBLES = [f"WFMOUTPRE:CH{ch};YMULT 1.5625E-3;YOFF -87.04;YZERO 0.0"
             for ch in SCAN_CHANNELS]
# adc_to_mv with those constants maps 0 → +136 mV and -120 → -51 mV, so a
# frame of zeros passes a 50 mV CH4 gate and a frame of -120 does not.
DC_THRESHOLD_MV = 50.0
CH4_PASS = bytes(SPF)
CH4_FAIL = bytes([256 - 120]) * SPF


def full_plan(num_angles=3, y_delta=0.05):
    """A small ROI, well inside the stage limits, with several rows per angle."""
    return build_plan(40.0, 30.0, 0.02, y_delta, num_angles, 0.01,
                      laser_freq_hz=LASER_FREQ_HZ, velocity_mm_s=VELOCITY_MM_S)


def bin_mhz(k: int) -> float:
    return k * SAMPLE_RATE / SPF / 1e6


def sine_frame(k: int) -> bytes:
    """One frame holding a pure sine at FFT bin `k`."""
    n = np.arange(SPF)
    return np.round(100 * np.sin(2 * math.pi * k * n / SPF)).astype(np.int8).tobytes()


def write_check(path, bins, n_masked_frames=0, plan=None, backgrounds=None):
    """A synthetic v11 file: angle `i`'s CH1 is a sine at FFT bin `bins[i]`.

    ``backgrounds`` supplies each angle's own background block; the default
    is a flat zero one per angle, which subtracts to nothing.
    """
    plan = plan if plan is not None else middle_row_plan(full_plan(len(bins)))
    f = create_scan_file(path, plan, SPF, SAMPLE_RATE, PREAMBLES,
                         version=VERSION_SAW_CHECK)
    try:
        for ai, pa in enumerate(plan.per_angle):
            write_background_block(
                f, bytes(SPF) if backgrounds is None else backgrounds[ai])
            wave = sine_frame(bins[ai])
            for ch in SCAN_CHANNELS:
                for fi in range(pa.n_frames):
                    if ch == 1:
                        f.write(wave)
                    elif ch == 3:
                        f.write(bytes(SPF))
                    else:
                        f.write(CH4_FAIL if fi < n_masked_frames else CH4_PASS)
    finally:
        f.close()
    return plan


# ── Plan reduction ───────────────────────────────────────────────────────────

def test_middle_row_plan_keeps_one_middle_row_per_angle():
    plan = full_plan(num_angles=3)
    check = middle_row_plan(plan)

    assert check.n_angles == plan.n_angles
    assert [pa.n_rows for pa in check.per_angle] == [1] * plan.n_angles
    for original, reduced in zip(plan.per_angle, check.per_angle, strict=True):
        mid = original.n_rows // 2
        assert reduced.y_positions == [original.y_positions[mid]]
        # The row is scanned exactly as the full scan would have scanned it.
        assert reduced.angle_deg == original.angle_deg
        assert reduced.x_start == original.x_start
        assert reduced.x_delta == original.x_delta
        assert reduced.n_frames == original.n_frames


def test_middle_row_plan_does_not_mutate_its_input():
    plan = full_plan(num_angles=3)
    before = [(pa.n_rows, list(pa.y_positions)) for pa in plan.per_angle]
    middle_row_plan(plan)
    assert [(pa.n_rows, pa.y_positions) for pa in plan.per_angle] == before


def test_every_angles_middle_row_crosses_the_roi_centre():
    """The premise the whole comparison rests on: one shared point on the sample."""
    plan = full_plan(num_angles=5)
    check = middle_row_plan(plan)
    cx = plan.x_start_nominal + plan.x_delta_nominal / 2
    cy = plan.y_start_nominal + plan.y_delta_nominal / 2
    for pa in check.per_angle:
        assert pa.x_start + pa.x_delta / 2 == pytest.approx(cx, abs=1e-6)
        # Within one row spacing — the middle row is a grid point, not exact.
        assert abs(pa.y_positions[0] - cy) <= plan.row_spacing


def test_middle_row_index_rule():
    assert [middle_row_index(n) for n in (1, 2, 3, 4, 6)] == [0, 1, 1, 2, 3]


def test_middle_row_plan_rejects_an_empty_plan():
    plan = full_plan(num_angles=1)
    plan.per_angle = []
    with pytest.raises(ScanGeometryError, match="no angles"):
        middle_row_plan(plan)


def test_middle_row_plan_rejects_an_angle_with_no_rows():
    plan = full_plan(num_angles=1)
    plan.per_angle[0].y_positions = []
    with pytest.raises(ScanGeometryError, match="no middle row"):
        middle_row_plan(plan)


# ── The v11 file ─────────────────────────────────────────────────────────────

def test_saw_check_write_read_roundtrip(tmp_path):
    out = tmp_path / "check.sras"
    plan = write_check(out, bins=(8, 8, 8))

    sras = SrasFile(out)
    assert sras.version == VERSION_SAW_CHECK
    assert sras.is_saw_check
    assert [s.status for s in sras.angle_status()] == ["OK"] * plan.n_angles
    assert [pa.n_rows for pa in sras.per_angle] == [1] * plan.n_angles
    sras.close()


def test_saw_check_rejects_a_multi_row_plan(tmp_path):
    plan = full_plan(num_angles=2)
    assert any(pa.n_rows > 1 for pa in plan.per_angle)
    with pytest.raises(ValueError, match="exactly one row per angle"):
        create_scan_file(tmp_path / "bad.sras", plan, SPF, SAMPLE_RATE,
                         PREAMBLES, version=VERSION_SAW_CHECK)
    assert not (tmp_path / "bad.sras").exists()


def test_unknown_version_rejected_at_write(tmp_path):
    with pytest.raises(ValueError, match="version 99"):
        create_scan_file(tmp_path / "bad.sras", middle_row_plan(full_plan(1)),
                         SPF, SAMPLE_RATE, PREAMBLES, version=99)


def test_a_scan_is_not_a_saw_check():
    """Legacy and current scans alike: only the check versions say check."""
    assert not SrasFile("tests/golden/complete.sras").is_saw_check
    assert VERSION not in (10, VERSION_SAW_CHECK)


# ── Acquisition through the engine ───────────────────────────────────────────

def run_engine(tmp_path, num_angles=3):
    trace = Trace()
    scope = FakeScope(trace, samples_per_frame=SPF)
    stage = FakeStage(trace, scope=scope)
    rotator = RotationAxis(FakeT3R(trace), RotationSettings())
    plan = full_plan(num_angles)
    check = middle_row_plan(plan)
    engine = ScanEngine(stage, scope, rotator, check, tmp_path / "check.sras",
                        callbacks=ScanCallbacks(),
                        file_version=VERSION_SAW_CHECK)
    return engine.run(), plan, check, trace


def test_engine_writes_a_complete_saw_check(tmp_path):
    result, plan, check, _ = run_engine(tmp_path)

    assert not result.aborted
    assert result.rows_written == check.n_angles      # exactly one row per angle
    assert result.angles_acquired == list(range(check.n_angles))

    sras = SrasFile(result.path)
    assert sras.is_saw_check
    assert [s.status for s in sras.angle_status()] == ["OK"] * check.n_angles
    assert [pa.y_positions for pa in sras.per_angle] == [
        [pytest.approx(original.y_positions[original.n_rows // 2], abs=1e-4)]
        for original in plan.per_angle
    ]
    sras.close()


def test_engine_visits_each_middle_row_once(tmp_path):
    _, _, check, trace = run_engine(tmp_path)
    y_moves = [round(c[2], 4) for c in trace.of("move_axis_absolute")
               if c[1] == 0x22]
    assert y_moves == [round(pa.y_positions[0], 4) for pa in check.per_angle]


def test_engine_still_writes_a_full_scan_by_default(tmp_path):
    trace = Trace()
    scope = FakeScope(trace, samples_per_frame=SPF)
    stage = FakeStage(trace, scope=scope)
    rotator = RotationAxis(FakeT3R(trace), RotationSettings())
    engine = ScanEngine(stage, scope, rotator, full_plan(1),
                        tmp_path / "scan.sras", callbacks=ScanCallbacks())
    result = engine.run()
    assert SrasFile(result.path).version == VERSION


# ── Analysis ─────────────────────────────────────────────────────────────────

def test_traces_report_the_injected_frequency(tmp_path):
    out = tmp_path / "check.sras"
    bins = (8, 9, 10)
    write_check(out, bins=bins)

    with SrasFile(out) as sras:
        traces = frequency_traces(sras, dc_threshold_mv=DC_THRESHOLD_MV)

    assert len(traces) == len(bins)
    for trace, k in zip(traces, bins, strict=True):
        assert np.allclose(trace.freq_mhz, bin_mhz(k))
        assert trace.median_mhz == pytest.approx(bin_mhz(k))
        assert trace.valid_fraction == 1.0
        assert trace.drift_mhz_per_mm == pytest.approx(0.0, abs=1e-6)


def test_background_subtraction_uses_each_angles_own(tmp_path):
    """Every angle is referenced against its own background, not angle 1's.

    Each angle's background here is a copy of that angle's own CH1 wave, so
    subtracting the right one leaves nothing to read at any angle — where
    reusing angle 1's everywhere would leave angles 2 and 3 reporting their
    sine unchanged.
    """
    out = tmp_path / "check.sras"
    bins = (8, 9, 10)
    backgrounds = [sine_frame(k) for k in bins]
    write_check(out, bins=bins, backgrounds=backgrounds)

    with SrasFile(out) as sras:
        assert sras.backgrounds == backgrounds
        plain = frequency_traces(sras, dc_threshold_mv=DC_THRESHOLD_MV)
        subtracted = frequency_traces(sras, dc_threshold_mv=DC_THRESHOLD_MV,
                                      subtract_background=True)

    assert [t.median_mhz for t in plain] == [pytest.approx(bin_mhz(k)) for k in bins]
    for trace in subtracted:
        assert np.isnan(trace.freq_mhz).all()


def test_masked_pixels_become_nan_not_zero(tmp_path):
    out = tmp_path / "check.sras"
    write_check(out, bins=(8, 8, 8), n_masked_frames=2)

    with SrasFile(out) as sras:
        traces = frequency_traces(sras, dc_threshold_mv=DC_THRESHOLD_MV)

    for trace in traces:
        assert np.isnan(trace.freq_mhz[:2]).all()
        assert np.isfinite(trace.freq_mhz[2:]).all()
        # A masked pixel must not drag the median toward 0 MHz.
        assert trace.median_mhz == pytest.approx(bin_mhz(8))
        assert trace.valid_fraction < 1.0


def test_traces_are_centred_on_a_common_offset(tmp_path):
    out = tmp_path / "check.sras"
    write_check(out, bins=(8, 9, 10))

    with SrasFile(out) as sras:
        traces = frequency_traces(sras, dc_threshold_mv=DC_THRESHOLD_MV)

    # Absolute X differs per angle (different bounding boxes); the offset the
    # viewer plots against does not, which is what puts the curves together.
    assert len({round(t.x_mm[0], 6) for t in traces}) > 1
    for trace in traces:
        assert trace.offset_mm[0] == pytest.approx(-trace.offset_mm[-1])


def test_angles_with_no_data_are_skipped(tmp_path):
    out = tmp_path / "check.sras"
    write_check(out, bins=(8, 8, 8))
    full = out.read_bytes()
    with SrasFile(out) as sras:
        last_offset = sras.angle_data_offset(2)
    out.write_bytes(full[:last_offset])          # angle 3 never acquired

    with SrasFile(out) as sras:
        traces = frequency_traces(sras, dc_threshold_mv=DC_THRESHOLD_MV)
    assert [t.angle_idx for t in traces] == [0, 1]


def test_summary_flags_agreeing_angles_as_good(tmp_path):
    out = tmp_path / "check.sras"
    write_check(out, bins=(8, 8, 8))

    with SrasFile(out) as sras:
        summary = alignment_summary(
            frequency_traces(sras, dc_threshold_mv=DC_THRESHOLD_MV))

    assert summary.n_angles == 3
    assert summary.median_mhz == pytest.approx(bin_mhz(8))
    assert summary.spread_mhz == pytest.approx(0.0)
    assert summary.spread_pct <= SPREAD_GOOD_PCT
    assert summary.level == "good"


def test_summary_flags_disagreeing_angles(tmp_path):
    out = tmp_path / "check.sras"
    write_check(out, bins=(8, 9, 10))

    with SrasFile(out) as sras:
        traces = frequency_traces(sras, dc_threshold_mv=DC_THRESHOLD_MV)
    summary = alignment_summary(traces)

    assert summary.spread_mhz == pytest.approx(bin_mhz(10) - bin_mhz(8))
    assert summary.level == "poor"
    assert summary.worst_angle_deg == traces[0].angle_deg   # lowest median
    assert summary.best_angle_deg == traces[2].angle_deg    # highest median
    assert f"{summary.spread_mhz:.3f} MHz" in summary.describe()


def test_summary_calls_out_a_mostly_masked_row(tmp_path):
    out = tmp_path / "check.sras"
    plan = middle_row_plan(full_plan(3))
    # Mask nearly every frame of every angle: the spread is meaningless then.
    write_check(out, bins=(8, 8, 8), plan=plan,
                n_masked_frames=max(pa.n_frames for pa in plan.per_angle) - 1)

    with SrasFile(out) as sras:
        summary = alignment_summary(
            frequency_traces(sras, dc_threshold_mv=DC_THRESHOLD_MV))

    assert summary.level == "poor"
    assert "DC threshold" in summary.describe()


def test_summary_of_nothing_is_not_a_crash():
    summary = alignment_summary([])
    assert summary.n_angles == 0 and summary.level == "poor"
    assert "No angle" in summary.describe()


def test_middle_row_of_a_full_v6_scan_is_readable():
    """The check's read-out applied to a finished scan, after the fact."""
    with SrasFile("tests/golden/complete.sras") as sras:
        traces = frequency_traces(sras, dc_threshold_mv=-1e6)
    assert len(traces) == sras.header.n_angles
    for trace, pa in zip(traces, sras.per_angle, strict=True):
        assert trace.row_idx == pa.n_rows // 2
        assert len(trace.freq_mhz) == pa.n_frames