scanengine-3/core/angle_inspect.py

"""Pre-scan angle inspection: park the rig on a point and let the operator look.

A multi-angle scan can take hours, and an angle that responds poorly produces
rows that look fine in the file but carry no usable SAW packet.  This drives
the rig through the same angles the scan will use, parking at a random point
inside each angle's own bounding box so the response can be judged on the
oscilloscope before committing to the run.

Headless and Qt-free, like ScanEngine: gui/inspect_bridge.py wraps it.

No waveform ever crosses this boundary.  The operator reads the scope screen
directly; this module's job is only to put the hardware in the right place and
the scope in a state worth looking at (see core.scope_inspect).
"""
from __future__ import annotations

import logging
import random
from dataclasses import dataclass
from typing import Callable

from core import scope_inspect
from core.rotation import RotationAxis
from core.scan_engine import (
    AXIS_X, AXIS_Y, SCAN_ACCEL_MM_S2, SCAN_VELOCITY_MM_S,
)
from core.scan_geometry import DEFAULT_STAGE_LIMITS, ScanPlan, StageLimits

logger = logging.getLogger(__name__)

# Positioning moves only — no data is taken while moving, so there is no
# reason to cross the tray at full scan velocity.
INSPECT_VELOCITY_MM_S = SCAN_VELOCITY_MM_S / 2.0


@dataclass(frozen=True)
class InspectionPoint:
    """Where the rig is parked, and which angle it is parked for."""
    angle_idx: int
    angle_deg: float
    x_mm: float
    y_mm: float

    def describe(self) -> str:
        return (f"Angle {self.angle_idx + 1} ({self.angle_deg:.1f}°)  "
                f"X={self.x_mm:.3f} mm  Y={self.y_mm:.3f} mm")


@dataclass
class InspectCallbacks:
    """Progress reporting.  Defaults are no-ops so the core needs no front end."""
    on_status: Callable[[str], None] = lambda msg: None
    on_point: Callable[[InspectionPoint], None] = lambda pt: None
    on_busy: Callable[[bool], None] = lambda busy: None


@dataclass
class _State:
    angle_idx: int = 0
    point: InspectionPoint | None = None
    started: bool = False
    rotator_ready: bool = False


class AngleInspector:
    """Drives stage + rotator to inspection points across a plan's angles."""

    def __init__(self, stage, scope, rotator: RotationAxis | None,
                 plan: ScanPlan,
                 callbacks: InspectCallbacks | None = None,
                 limits: StageLimits = DEFAULT_STAGE_LIMITS,
                 rng: random.Random | None = None):
        self._stage = stage
        self._scope = scope
        self._rotator = rotator
        self._plan = plan
        self._cb = callbacks if callbacks is not None else InspectCallbacks()
        self._limits = limits
        # Injectable so tests can pin the point selection.
        self._rng = rng if rng is not None else random.Random()
        self._st = _State()

    # ── Introspection ─────────────────────────────────────────────────────────

    @property
    def n_angles(self) -> int:
        return self._plan.n_angles

    @property
    def angle_idx(self) -> int:
        return self._st.angle_idx

    @property
    def current_point(self) -> InspectionPoint | None:
        return self._st.point

    def angle_labels(self) -> list[str]:
        return [f"Angle {i + 1}/{self.n_angles} — {pa.angle_deg:.2f}°"
                for i, pa in enumerate(self._plan.per_angle)]

    # ── Lifecycle ─────────────────────────────────────────────────────────────

    def start(self) -> InspectionPoint:
        """Configure the hardware and park on the first angle."""
        if self._stage is None:
            raise RuntimeError("BBD202 not connected")
        if self._scope is None:
            raise RuntimeError("Oscilloscope not connected")

        self._st.rotator_ready = (self._rotator is not None
                                  and self._rotator.is_available)
        if self.n_angles > 1 and not self._st.rotator_ready:
            raise RuntimeError(
                f"Inspecting {self.n_angles} angles requires the T3R rotation "
                "stage (GR-axis), but it is not connected. Connect T3R from "
                "the T3R panel, or inspect a single-angle plan."
            )

        self._cb.on_busy(True)
        try:
            self._cb.on_status("Configuring stage for inspection …")
            ctrl = self._stage
            for axis in (AXIS_X, AXIS_Y):
                ctrl.set_velocity_params(axis,
                                         max_velocity=INSPECT_VELOCITY_MM_S,
                                         acceleration=SCAN_ACCEL_MM_S2)
            # Nothing here is gated, and an armed trigger output would keep
            # driving the gate line on every positioning move.
            ctrl.set_trigger_gate_off(AXIS_X)

            if self._st.rotator_ready:
                self._cb.on_status("Configuring GR axis …")
                self._rotator.configure()

            self._cb.on_status("Configuring oscilloscope for inspection …")
            scope_inspect.configure_inspection(self._scope)

            self._st.started = True
            return self._goto(0, new_point=True)
        finally:
            self._cb.on_busy(False)

    def stop(self) -> None:
        """Stop the sweep and send the rotator home.  Safe to call twice."""
        if not self._st.started:
            return
        self._st.started = False
        self._cb.on_busy(True)
        try:
            try:
                scope_inspect.stop_inspection(self._scope)
            except Exception:
                logger.exception("Could not stop the inspection acquisition")
            if self._st.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")
            self._cb.on_status("Inspection finished.")
        finally:
            self._cb.on_busy(False)

    # ── Navigation ────────────────────────────────────────────────────────────

    def goto_angle(self, angle_idx: int) -> InspectionPoint:
        """Rotate to `angle_idx` and park on a fresh random point there."""
        self._require_started()
        self._cb.on_busy(True)
        try:
            return self._goto(angle_idx, new_point=True)
        finally:
            self._cb.on_busy(False)

    def next_angle(self) -> InspectionPoint:
        """Advance one angle, wrapping at the end."""
        return self.goto_angle((self._st.angle_idx + 1) % self.n_angles)

    def prev_angle(self) -> InspectionPoint:
        return self.goto_angle((self._st.angle_idx - 1) % self.n_angles)

    def new_point(self) -> InspectionPoint:
        """Re-roll the point within the current angle, without rotating.

        One point can be unrepresentative — a bad spot on the sample looks the
        same as a bad angle.  Re-rolling a few times is how you tell them
        apart, so this deliberately skips the rotation.
        """
        self._require_started()
        self._cb.on_busy(True)
        try:
            return self._goto(self._st.angle_idx, new_point=True, rotate=False)
        finally:
            self._cb.on_busy(False)

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

    def _require_started(self):
        if not self._st.started:
            raise RuntimeError("Inspection has not been started")

    def _goto(self, angle_idx: int, new_point: bool,
              rotate: bool = True) -> InspectionPoint:
        if not 0 <= angle_idx < self.n_angles:
            raise IndexError(
                f"Angle {angle_idx} out of range (plan has {self.n_angles})")

        pa = self._plan.per_angle[angle_idx]
        self._st.angle_idx = angle_idx

        if rotate and self._st.rotator_ready:
            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)

        point = self._pick_point(angle_idx) if new_point else self._st.point

        self._cb.on_status(f"Moving to {point.describe()} …")
        # Y first, then X — the same order the scan uses to reach a row.
        self._stage.move_axis_absolute(AXIS_Y, point.y_mm, timeout=60.0)
        self._stage.move_axis_absolute(AXIS_X, point.x_mm, timeout=60.0)

        self._st.point = point
        self._cb.on_point(point)
        self._cb.on_status(f"Parked at {point.describe()}")
        return point

    def _pick_point(self, angle_idx: int) -> InspectionPoint:
        """A random point on this angle's scan grid.

        Y is drawn from the angle's actual row positions and X uniformly from
        its data window, so the point is somewhere the scan would really
        sample — not merely inside the bounding box.
        """
        pa = self._plan.per_angle[angle_idx]
        if not pa.y_positions:
            raise ValueError(f"Angle {angle_idx + 1} has no rows to inspect")

        y = self._rng.choice(pa.y_positions)
        x = self._rng.uniform(pa.x_start, pa.x_start + pa.x_delta)

        lim = self._limits
        if not (lim.x_min <= x <= lim.x_max and lim.y_min <= y <= lim.y_max):
            raise ValueError(
                f"Inspection point X={x:.3f} Y={y:.3f} is outside the stage "
                f"travel ({lim.x_min}–{lim.x_max} × {lim.y_min}–{lim.y_max} mm)"
            )
        return InspectionPoint(angle_idx=angle_idx, angle_deg=pa.angle_deg,
                               x_mm=x, y_mm=y)