scanengine-3/core/scan_engine.py

"""Headless SRAS scan engine.

Takes plain hardware drivers, a ScanPlan, and callbacks — no Qt, no
widgets.  ``run()`` blocks, so the caller owns the thread; GUIs wrap this
with gui.scan_bridge.QtScanController, which adapts the callbacks to Qt
signals.  A CLI or a simpler GUI can drive the same engine with nothing but
functions.
"""
from __future__ import annotations

import logging
import threading
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable

from core import scope_burst, scope_sras
from core.rotation import RotationAxis
from core.scan_geometry import ScanPlan, validate_plan
from core.sras_format import SCAN_CHANNELS, create_scan_file

logger = logging.getLogger(__name__)

SCAN_VELOCITY_MM_S = 100.0
SCAN_ACCEL_MM_S2 = 1500.0
LASER_FREQ_HZ = 20000.0    # laser pulse frequency during data acquisition
# Theoretical ramp distance: d = v² / (2a) = 100² / (2×1500) ≈ 3.33 mm
SCAN_RAMP_MM = SCAN_VELOCITY_MM_S**2 / (2.0 * SCAN_ACCEL_MM_S2)
# Extra buffer added to both ends of the ramp.  The BBD202 controller begins
# decelerating slightly before the theoretical point to avoid overshoot,
# which drops TRIGOUT_MAXV early and clips the last few data points.
SCAN_RAMP_BUFFER_MM = 1.0

AXIS_X = 0x21
AXIS_Y = 0x22


class ScanAborted(Exception):
    """Raised inside the engine thread to unwind a scan cleanly."""


@dataclass
class ResumeTarget:
    """One angle selected for (re)acquisition in an existing file."""
    angle_idx: int
    data_offset: int
    n_rows: int
    angle_deg: float


@dataclass
class ResumeState:
    path: Path
    targets: list[ResumeTarget]
    samples_per_frame: int

    @property
    def target_indices(self) -> set[int]:
        return {t.angle_idx for t in self.targets}


@dataclass
class ScanCallbacks:
    """Progress reporting hooks.  Every one is optional."""
    on_status: Callable[[str], None] = lambda msg: None
    on_started: Callable[[], None] = lambda: None
    on_row_started: Callable[[int, int, int, int], None] = lambda r, nr, a, na: None
    on_row_done: Callable[[int, int, int, int], None] = lambda r, nr, a, na: None
    on_dc_bias: Callable[[int, list], None] = lambda row, means: None
    on_paused_changed: Callable[[bool], None] = lambda paused: None
    # Blocking operator prompt: must not return until acknowledged.
    prompt: Callable[[str, str], None] = lambda title, msg: None


@dataclass
class ScanResult:
    path: Path
    rows_written: int = 0
    aborted: bool = False
    angles_acquired: list[int] = field(default_factory=list)


class ScanEngine:
    """Runs a full SRAS acquisition: stage, rotation, scope, and file output.

    Constructed with the concrete drivers (not worker/queue wrappers), so
    any front end can reuse it::

        engine = ScanEngine(stage, scope, rotator, plan, out_path,
                            callbacks=ScanCallbacks(on_status=print))
        result = engine.run()          # blocking
    """

    def __init__(self, stage, scope, rotator: RotationAxis | None,
                 plan: ScanPlan, out_path: Path,
                 resume: ResumeState | None = None,
                 callbacks: ScanCallbacks | None = None,
                 burst_mode: bool = False):
        self._stage = stage
        self._scope = scope
        self._rotator = rotator
        self._plan = plan
        self._out_path = Path(out_path)
        self._resume = resume
        self._cb = callbacks if callbacks is not None else ScanCallbacks()
        # Burst mode acquires as many whole rows per FastFrame acquisition as
        # the scope's frame memory holds, instead of one row per acquisition.
        self._burst_mode = burst_mode
        self._max_frames = 0
        self._preflight_done = False

        self._abort = threading.Event()
        self._resume_event = threading.Event()
        self._resume_event.set()   # set = running, cleared = pause requested

    # ── External control (thread-safe) ────────────────────────────────────────

    def abort(self):
        self._abort.set()
        self._resume_event.set()   # unblock a paused scan so it can exit

    def pause(self):
        """Request a pause; takes effect at the next row boundary."""
        self._resume_event.clear()

    def resume(self):
        self._resume_event.set()

    @property
    def aborted(self) -> bool:
        return self._abort.is_set()

    # ── Internals ─────────────────────────────────────────────────────────────

    def _check_abort(self):
        if self._abort.is_set():
            raise ScanAborted("Scan aborted by user.")

    def _pause_point(self):
        """Block here (between rows, hardware idle) while a pause is requested."""
        if self._resume_event.is_set():
            self._check_abort()
            return
        self._cb.on_status(
            "Scan paused — lasers may be switched off. "
            "Turn lasers back on before resuming."
        )
        self._cb.on_paused_changed(True)
        while not self._resume_event.wait(0.2):
            if self._abort.is_set():
                break
        self._cb.on_paused_changed(False)
        self._check_abort()
        self._cb.on_status("Scan resumed.")

    def _prompt(self, title: str, message: str):
        self._cb.prompt(title, message)
        self._check_abort()

    # ── Main sequence ─────────────────────────────────────────────────────────

    def run(self) -> ScanResult:
        """Execute the scan.  Blocking; returns a ScanResult.

        Raises ScanGeometryError for an unrunnable plan, RuntimeError for
        missing/mismatched hardware, or ScanAborted if the operator aborts.
        """
        plan = self._plan
        per_angle = plan.per_angle
        n_angles = plan.n_angles
        result = ScanResult(path=self._out_path)

        validate_plan(plan, SCAN_RAMP_MM, SCAN_RAMP_BUFFER_MM)

        geometry_summary = ", ".join(
            f"{pa.angle_deg:.1f}°: {pa.n_rows} row(s) × {pa.n_frames} pts/row"
            for pa in per_angle
        )
        self._cb.on_status(
            f"Scan geometry: {n_angles} angle(s), {plan.total_rows} row(s) total "
            f"(per-angle bounding box)  |  save → {self._out_path.parent}\n"
            f"{geometry_summary}"
        )
        self._cb.on_started()

        if self._stage is None:
            raise RuntimeError("BBD202 not connected")
        if self._scope is None:
            raise RuntimeError("Oscilloscope not connected")
        rotator_ready = self._rotator is not None and self._rotator.is_available
        if n_angles > 1 and not rotator_ready:
            raise RuntimeError(
                f"NumAngles={n_angles} requires the T3R rotation stage (GR-axis), "
                "but it is not connected. Connect T3R from the T3R panel before "
                "starting a multi-angle scan, or set NumAngles to 1."
            )

        if rotator_ready:
            s = self._rotator.settings
            self._cb.on_status(
                f"Configuring GR axis: {s.microsteps} µsteps, "
                f"{s.run_current_ma}/{s.hold_current_ma} mA run/hold …"
            )
            self._rotator.configure()
            time.sleep(0.2)

        self._prepare_stage()
        samples_per_frame = self._prepare_scope()
        scan_file = self._open_output(samples_per_frame, result)

        try:
            self._scan_loop(scan_file, samples_per_frame, result)
        finally:
            scan_file.close()
            # Leave the X trigger output inactive.  Burst mode toggles it every
            # row and could exit from either state; the per-row path used to
            # leave TRIGOUT_MAXV armed for the rest of the session, which keeps
            # driving the gate line on every later jog.
            try:
                self._stage.set_trigger_gate_off(AXIS_X)
            except Exception:
                logger.exception("Could not return the X trigger output to idle")
            # Return the GR axis home regardless of abort or error
            if rotator_ready and abs(self._rotator.current_deg) > 0.001:
                self._cb.on_status("Returning GR to home …")
                try:
                    self._rotator.return_to_zero()
                except Exception:
                    logger.exception("GR return-to-home failed")

        if self._abort.is_set():
            result.aborted = True
            raise ScanAborted("Scan aborted by user.")
        self._cb.on_status("Scan complete.")
        return result

    def _prepare_stage(self):
        ctrl = self._stage
        self._cb.on_status("Enabling stage axes …")
        if not ctrl.am_enabled[0]:
            ctrl.enable_axis(AXIS_X)
        if not ctrl.am_enabled[1]:
            ctrl.enable_axis(AXIS_Y)
        time.sleep(0.2)

        if not ctrl.am_homed[0] or not ctrl.am_homed[1]:
            self._cb.on_status("Homing stage (may take up to 2 min) …")
            if not ctrl.am_homed[0]:
                ctrl.home_axis(AXIS_X, timeout=120.0)
            if not ctrl.am_homed[1]:
                ctrl.home_axis(AXIS_Y, timeout=120.0)

        self._cb.on_status("Setting scan velocity …")
        ctrl.set_velocity_params(AXIS_X, max_velocity=SCAN_VELOCITY_MM_S,
                                 acceleration=SCAN_ACCEL_MM_S2)
        ctrl.set_velocity_params(AXIS_Y, max_velocity=SCAN_VELOCITY_MM_S,
                                 acceleration=SCAN_ACCEL_MM_S2)
        # X trigger: logic-high output while the stage is at maximum velocity.
        # Burst mode arms it per acquiring pass instead — a burst spans several
        # rows with the scope running throughout, so leaving it armed would let
        # the flyback trigger frames between rows.
        if self._burst_mode:
            ctrl.arm_scan_gate(AXIS_X, False)
        else:
            ctrl.set_trigger_trigout_maxv(AXIS_X)

    def _prepare_scope(self) -> int:
        self._cb.on_status("Configuring oscilloscope …")
        samples_per_frame = scope_sras.configure_acquisition(self._scope)

        if self._resume is None:
            self._preambles = scope_sras.read_preambles(self._scope, SCAN_CHANNELS)
            self._prompt(
                "Background Capture",
                "Please ensure the Helios laser is ON and the Genesis laser is OFF,\n"
                "then click OK to capture the background waveform."
            )
            self._background = scope_sras.capture_background(
                self._scope, should_abort=self._abort.is_set,
                on_status=self._cb.on_status)
            self._prompt(
                "Begin Scanning",
                "Background captured successfully.\n\n"
                "Please ensure the Genesis laser is back ON,\n"
                "then click OK to begin scanning."
            )
        else:
            # Resuming: the file's existing background waveform and channel
            # preambles are reused as-is (the format has no way to replace
            # them without rewriting the whole file), so background capture is
            # skipped.  Sanity-check that this scope still produces the same
            # record length the file was started with — a mismatch would
            # silently corrupt the ragged per-row byte layout on append.
            if samples_per_frame != self._resume.samples_per_frame:
                raise RuntimeError(
                    f"Oscilloscope record length ({samples_per_frame} samples/frame) "
                    f"does not match the {self._resume.samples_per_frame} samples/frame "
                    f"this scan file was started with — cannot safely resume."
                )
            targets = ", ".join(str(t.angle_idx + 1) for t in self._resume.targets)
            self._prompt(
                "Resume Scan",
                f"Resuming {self._resume.path.name} — will (re)acquire "
                f"angle(s) {targets} of {self._plan.n_angles}.\n\n"
                "Please re-home the GR axis to 0° before continuing — the scan "
                "will rotate it directly from angle to angle before scanning resumes.\n\n"
                "Please ensure the Genesis laser is ON,\n"
                "then click OK to continue scanning."
            )

        scope_sras.configure_scan_trigger(self._scope)

        if self._burst_mode:
            # Horizontal settings are fixed by now, so the capacity is stable
            # for the whole scan; only rows-per-burst varies (n_frames is
            # per-angle).
            self._max_frames = scope_burst.max_frames(self._scope)
            self._cb.on_status(
                f"Burst mode: scope holds {self._max_frames} frames "
                f"({samples_per_frame} samples/frame)"
            )
        return samples_per_frame

    def _open_output(self, samples_per_frame: int, result: ScanResult):
        if self._resume is not None:
            result.path = self._resume.path
            self._cb.on_status(
                f"Resuming {self._resume.path.name} — "
                f"{len(self._resume.targets)} angle(s) to (re)acquire …"
            )
            return open(self._resume.path, "r+b")
        return create_scan_file(
            self._out_path, self._plan, samples_per_frame,
            scope_sras.SAMPLE_RATE_HZ, self._preambles, self._background,
        )

    def _scan_loop(self, scan_file, samples_per_frame: int, result: ScanResult):
        plan = self._plan
        n_angles = plan.n_angles
        x_ramp_total = SCAN_RAMP_MM + SCAN_RAMP_BUFFER_MM
        scope = self._scope

        targets_by_ai = None
        if self._resume is not None:
            targets_by_ai = {t.angle_idx: t for t in self._resume.targets}

        # Both fresh and resumed scans assume the GR axis starts at home (0°)
        # — the resume prompt instructs the operator to re-home it — so the
        # first move always rotates directly from 0° to the starting angle.
        for ai, pa in enumerate(plan.per_angle):
            if targets_by_ai is not None and ai not in targets_by_ai:
                continue    # not selected for (re)acquisition
            self._pause_point()

            if targets_by_ai is not None:
                # Interior angles may already have valid data on either side,
                # so seek to this angle's fixed offset rather than relying on
                # the file's current position.
                scan_file.seek(targets_by_ai[ai].data_offset)

            if self._rotator is not None and self._rotator.is_available:
                delta = pa.angle_deg - self._rotator.current_deg
                if abs(delta) > 0.001:
                    self._cb.on_status(
                        f"Rotating GR to {pa.angle_deg:.1f}° (Δ{delta:+.1f}°) …")
                    self._rotator.rotate_to(pa.angle_deg)

            if self._burst_mode:
                # Burst mode sizes the FastFrame count from the scope's whole
                # capacity instead (see scope_burst.start_burst), so there is
                # nothing to re-arm per angle here.
                self._scan_rows_burst(scan_file, pa, ai, n_angles,
                                      samples_per_frame, result, x_ramp_total)
            else:
                # Each angle's bounding box gives it its own points/row count,
                # so the scope's FastFrame count must be re-armed per angle.
                scope.set_fastframe_count(pa.n_frames)
                self._scan_rows_serial(scan_file, pa, ai, n_angles,
                                       samples_per_frame, result, x_ramp_total)

            result.angles_acquired.append(ai)

    # ── Per-row acquisition (one FastFrame acquisition per row) ───────────────

    def _scan_rows_serial(self, scan_file, pa, ai: int, n_angles: int,
                          samples_per_frame: int, result: ScanResult,
                          x_ramp_total: float):
        ctrl = self._stage
        scope = self._scope

        for ri, y_pos in enumerate(pa.y_positions):
            self._pause_point()

            self._cb.on_row_started(ri + 1, pa.n_rows, ai + 1, n_angles)
            self._cb.on_status(
                f"Angle {ai+1}/{n_angles}  Row {ri+1}/{pa.n_rows}  "
                f"(Y={y_pos:.3f} mm)"
            )

            # Position the stage one ramp-length + buffer before the data
            # window so it is at full velocity before x_start.
            ctrl.move_axis_absolute(AXIS_Y, y_pos, timeout=60.0)
            ctrl.move_axis_absolute(AXIS_X, pa.x_start - x_ramp_total, timeout=30.0)

            scope_sras.arm_row(scope)

            # Data window + ramp + buffer run-off, so the stage does not
            # begin decelerating before the last point.
            x_end = pa.x_start + pa.x_delta + x_ramp_total
            ctrl.move_axis_absolute(AXIS_X, x_end, timeout=120.0)

            scope_sras.finish_row(scope)
            self._write_row(scan_file, samples_per_frame, ri, pa.n_frames)

            result.rows_written += 1
            self._cb.on_row_done(ri + 1, pa.n_rows, ai + 1, n_angles)

    def _write_row(self, scan_file, samples_per_frame: int, row_idx: int,
                   n_frames: int):
        """Stream every channel from the scope into the file.

        CH3 is the max-vel gate signal — no useful waveform data — so zeroed
        frames are written to keep the file layout intact.
        """
        scope = self._scope
        ch_bytes = n_frames * samples_per_frame
        for ch in SCAN_CHANNELS:
            if ch == 3:
                self._cb.on_status("Writing zeroed CH3 frames …")
                scan_file.write(bytes(ch_bytes))
                continue

            self._cb.on_status(f"Fetching CH{ch} data …")
            waveforms = scope_sras.transfer_channel(scope, ch)
            if ch == SCAN_CHANNELS[0]:
                self._warn_frame_delta(row_idx, len(waveforms), n_frames)
            row = scope_burst.normalize_row(
                b"".join(waveforms), 0, len(waveforms), n_frames, samples_per_frame)
            if ch == 4:
                self._cb.on_dc_bias(row_idx + 1, scope_burst.frame_means_block(
                    row, 0, n_frames, samples_per_frame))
            scan_file.write(row)

    # ── Burst acquisition (many whole rows per FastFrame acquisition) ─────────

    def _scan_rows_burst(self, scan_file, pa, ai: int, n_angles: int,
                         samples_per_frame: int, result: ScanResult,
                         x_ramp_total: float):
        """Acquire the angle in bursts of as many whole rows as the scope holds.

        One ACQuire:STATE RUN spans the whole burst, so the gate is armed only
        for each acquiring pass and dropped for the flyback — otherwise the
        return move would reach max velocity and inject frames between rows.
        """
        scope = self._scope
        n_frames = pa.n_frames
        x_lead_in = pa.x_start - x_ramp_total
        x_end = pa.x_start + pa.x_delta + x_ramp_total

        if not self._preflight_done:
            # Once per scan: the gate wiring can't change between angles, and
            # the check costs two row-times.
            self._gate_off_preflight(x_lead_in, x_end)
            self._preflight_done = True

        row = 0
        while row < pa.n_rows:
            self._pause_point()
            n_burst = scope_burst.rows_per_burst(
                self._max_frames, n_frames, samples_per_frame, pa.n_rows - row)
            self._cb.on_status(
                f"Angle {ai+1}/{n_angles}  Rows {row+1}-{row+n_burst}/{pa.n_rows} "
                f"in one acquisition ({n_burst * n_frames} frames) …"
            )

            burst_start = scan_file.tell()
            cumulative = []
            baseline = scope_burst.start_burst(scope, self._max_frames)
            try:
                for r in range(n_burst):
                    self._check_abort()
                    self._cb.on_row_started(row + r + 1, pa.n_rows,
                                            ai + 1, n_angles)
                    self._acquire_gated_row(pa.y_positions[row + r],
                                            x_lead_in, x_end)
                    total = scope_burst.frames_acquired(scope)
                    if total >= self._max_frames:
                        raise RuntimeError(
                            f"FastFrame buffer full ({total}/{self._max_frames} "
                            f"frames) at row {row + r + 1} — later rows in this "
                            "burst would be misattributed.  Raise "
                            "scope_burst.BURST_FRAME_HEADROOM and rerun."
                        )
                    cumulative.append(total - baseline)
            finally:
                scope_burst.stop_burst(scope)

            counts = scope_burst.split_row_counts(cumulative)
            self._write_burst(scan_file, burst_start, row, counts,
                              n_frames, samples_per_frame)

            for r in range(n_burst):
                result.rows_written += 1
                self._cb.on_row_done(row + r + 1, pa.n_rows, ai + 1, n_angles)
            row += n_burst

    def _acquire_gated_row(self, y_pos: float, x_lead_in: float, x_end: float):
        """One row: step Y, fly back gated off, then acquire on the +X pass."""
        ctrl = self._stage
        ctrl.move_axis_absolute(AXIS_Y, y_pos, timeout=60.0)
        ctrl.move_axis_absolute(AXIS_X, x_lead_in, timeout=30.0)
        ctrl.arm_scan_gate(AXIS_X, True)
        ctrl.move_axis_absolute(AXIS_X, x_end, timeout=120.0)
        ctrl.arm_scan_gate(AXIS_X, False)
        time.sleep(scope_burst.BURST_ROW_SETTLE_S)

    def _gate_off_preflight(self, x_lead_in: float, x_end: float):
        """Prove the gate really gates before trusting a multi-row burst.

        The value that makes the BBD trigger output idle low is not settled by
        the protocol docs (see apt_constants.TRIGOUT_GATE_OFF), and getting it
        wrong fills every burst with flyback frames that silently shift the
        file.  The scope already measures the gate on CH3, so this needs no
        bench probe: one gated-off flyback must acquire nothing, and one gated
        pass must acquire something — the second half is what stops a dark
        laser from making the first half pass vacuously.

        Leaves the stage parked at x_end, where the burst loop expects it.
        """
        ctrl, scope = self._stage, self._scope
        self._cb.on_status("Burst preflight: checking the stage gate …")

        ctrl.arm_scan_gate(AXIS_X, False)
        ctrl.move_axis_absolute(AXIS_X, x_end, timeout=120.0)
        baseline = scope_burst.start_burst(scope, self._max_frames)
        ctrl.move_axis_absolute(AXIS_X, x_lead_in, timeout=120.0)
        scope_burst.stop_burst(scope)
        leaked = scope_burst.frames_acquired(scope) - baseline

        ctrl.arm_scan_gate(AXIS_X, True)
        baseline = scope_burst.start_burst(scope, self._max_frames)
        ctrl.move_axis_absolute(AXIS_X, x_end, timeout=120.0)
        ctrl.arm_scan_gate(AXIS_X, False)
        scope_burst.stop_burst(scope)
        gated = scope_burst.frames_acquired(scope) - baseline

        if gated <= 0:
            raise RuntimeError(
                "Burst preflight: no frames acquired with the gate armed. "
                "Check that the Genesis laser is pulsing (CH2) and that the "
                "BBD X trigger output reaches CH3 before scanning."
            )
        if leaked:
            raise RuntimeError(
                f"Burst preflight: {leaked} frame(s) acquired during a flyback "
                "that should have been gated off — the BBD trigger output is "
                "not idling low.  Set apt_constants.TRIGOUT_GATE_OFF to "
                "TriggerBitsServo.TRIGOUT_HIGH and retry, or use per-row "
                "acquisition."
            )
        self._cb.on_status(
            f"Burst preflight OK ({gated} frames gated on, 0 leaked).")

    def _write_burst(self, scan_file, burst_start: int, first_row: int,
                     counts: list[int], n_frames: int, samples_per_frame: int):
        """Deinterleave one burst into the file's per-row, per-channel blocks.

        The wire is channel-major (every row of CH1, then every row of CH4);
        the file is row-major with channels inner.  Writing one channel at a
        time to strided offsets keeps peak memory at a single channel's burst
        instead of the whole thing.
        """
        scope = self._scope
        ch_bytes = n_frames * samples_per_frame
        row_bytes = len(SCAN_CHANNELS) * ch_bytes
        total_frames = sum(counts)

        for r, count in enumerate(counts):
            self._warn_frame_delta(first_row + r, count, n_frames)

        for ch_idx, ch in enumerate(SCAN_CHANNELS):
            if ch == 3:
                self._cb.on_status("Writing zeroed CH3 frames …")
                blob = None
            else:
                self._cb.on_status(
                    f"Fetching CH{ch} burst ({total_frames} frames) …")
                blob = scope_burst.transfer_burst(scope, ch, total_frames,
                                                  samples_per_frame)
            src = 0
            zeros = bytes(ch_bytes) if blob is None else None
            for r, count in enumerate(counts):
                scan_file.seek(burst_start + r * row_bytes + ch_idx * ch_bytes)
                if blob is None:
                    scan_file.write(zeros)
                else:
                    row = scope_burst.normalize_row(
                        blob, src, count, n_frames, samples_per_frame)
                    if ch == 4:
                        self._cb.on_dc_bias(
                            first_row + r + 1,
                            scope_burst.frame_means_block(
                                row, 0, n_frames, samples_per_frame))
                    scan_file.write(row)
                src += count * samples_per_frame
            del blob

        scan_file.seek(burst_start + len(counts) * row_bytes)

    def _warn_frame_delta(self, row_idx: int, count: int, n_frames: int):
        if count == n_frames:
            return
        verb = "zero-padded" if count < n_frames else "truncated"
        msg = (f"Row {row_idx + 1}: {count} frames acquired, {n_frames} "
               f"expected — {verb} to keep the file layout intact.")
        logger.warning(msg)
        self._cb.on_status(msg)