scanengine-3/core/scan_geometry.py

"""Scan geometry planning: angle sequences, rotated bounding boxes, travel
limits, and ETA math.  Pure Python — no Qt, no hardware.
"""
from __future__ import annotations

import math
import time
from collections import deque
from dataclasses import dataclass, field


class ScanGeometryError(ValueError):
    """Scan geometry that cannot be executed (bad inputs or off-stage)."""


@dataclass(frozen=True)
class StageLimits:
    """Usable travel of the scanning stage in mm (MLS203-1)."""
    x_min: float = 0.0
    x_max: float = 110.0
    y_min: float = 0.0
    y_max: float = 75.0


DEFAULT_STAGE_LIMITS = StageLimits()


@dataclass
class AngleGeometry:
    """One rotation angle's scan extent — also the on-disk v6 geometry row."""
    angle_deg: float
    x_start: float
    x_delta: float
    n_frames: int
    n_rows: int
    y_positions: list[float] = field(default_factory=list)


@dataclass
class ScanPlan:
    x_start_nominal: float
    y_start_nominal: float
    x_delta_nominal: float
    y_delta_nominal: float
    row_spacing: float
    velocity_mm_s: float
    laser_freq_hz: float
    per_angle: list[AngleGeometry] = field(default_factory=list)

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

    @property
    def angles(self) -> list[float]:
        return [pa.angle_deg for pa in self.per_angle]

    @property
    def total_rows(self) -> int:
        return sum(pa.n_rows for pa in self.per_angle)


def build_plan(x_start: float, y_start: float, x_delta: float, y_delta: float,
               num_angles: int, row_spacing: float, *,
               laser_freq_hz: float, velocity_mm_s: float,
               rotation_sign: int = -1) -> ScanPlan:
    """Compute the per-angle scan geometry for a nominal ROI.

    Each angle only needs to physically scan the bounding box of the nominal
    (x_start, y_start, x_delta, y_delta) rectangle rotated by THAT angle --
    not the worst case across all angles -- so the X extent (and therefore
    points/row) and the row count are computed per angle.
    """
    if x_delta <= 0:
        raise ScanGeometryError("XD must be > 0")
    if row_spacing <= 0:
        raise ScanGeometryError("RowSpacing must be > 0")
    if num_angles < 1:
        raise ScanGeometryError("NumAngles must be ≥ 1")

    # Signed so the recorded/commanded angle sequence reflects the GR
    # stage's actual physical rotation direction.
    if num_angles > 1:
        angles = [rotation_sign * i * 180.0 / (num_angles - 1) for i in range(num_angles)]
    else:
        angles = [0.0]

    cx = x_start + x_delta / 2.0
    cy = y_start + y_delta / 2.0
    per_angle = []
    for a in angles:
        r = math.radians(a)
        bb_w = abs(x_delta * math.cos(r)) + abs(y_delta * math.sin(r))
        bb_h = abs(x_delta * math.sin(r)) + abs(y_delta * math.cos(r))
        a_x_start = cx - bb_w / 2.0
        a_y_start = cy - bb_h / 2.0
        a_n_rows = max(1, round(bb_h / row_spacing) + 1) if bb_h > 0 else 1
        a_n_frames = max(1, round(bb_w * laser_freq_hz / velocity_mm_s))
        per_angle.append(AngleGeometry(
            angle_deg=a,
            x_start=a_x_start,
            x_delta=bb_w,
            n_frames=a_n_frames,
            n_rows=a_n_rows,
            y_positions=[a_y_start + i * row_spacing for i in range(a_n_rows)],
        ))

    return ScanPlan(
        x_start_nominal=x_start, y_start_nominal=y_start,
        x_delta_nominal=x_delta, y_delta_nominal=y_delta,
        row_spacing=row_spacing,
        velocity_mm_s=velocity_mm_s, laser_freq_hz=laser_freq_hz,
        per_angle=per_angle,
    )


def validate_plan(plan: ScanPlan, ramp_mm: float, ramp_buffer_mm: float,
                  limits: StageLimits = DEFAULT_STAGE_LIMITS) -> None:
    """Raise ScanGeometryError if any angle's physical move leaves the stage.

    The actual X move starts one ramp-length + buffer before x_start and ends
    one ramp-length + buffer after x_start + x_delta, so the stage is at full
    velocity across the whole data window.
    """
    x_ramp_total = ramp_mm + ramp_buffer_mm
    for pa in plan.per_angle:
        x_move_start = pa.x_start - x_ramp_total
        x_move_end = pa.x_start + pa.x_delta + x_ramp_total
        if x_move_start < limits.x_min:
            raise ScanGeometryError(
                f"Angle {pa.angle_deg:.1f}°: scan pre-ramp start ({x_move_start:.3f} mm) "
                f"is below the X axis minimum ({limits.x_min:g} mm). Reduce XD/YD or move XS/YS "
                f"so every rotation angle's bounding box stays on-stage "
                f"(SCAN_RAMP_MM={ramp_mm:.3f} + SCAN_RAMP_BUFFER_MM={ramp_buffer_mm:.3f})."
            )
        if x_move_end > limits.x_max:
            raise ScanGeometryError(
                f"Angle {pa.angle_deg:.1f}°: scan run-off end ({x_move_end:.3f} mm) "
                f"exceeds the X axis maximum ({limits.x_max:g} mm). Reduce XD/YD or move XS/YS "
                f"so every rotation angle's bounding box stays on-stage."
            )
        y_min = min(pa.y_positions)
        y_max = max(pa.y_positions)
        if y_min < limits.y_min:
            raise ScanGeometryError(
                f"Angle {pa.angle_deg:.1f}°: scan Y range starts at {y_min:.3f} mm, "
                f"below the Y axis minimum ({limits.y_min:g} mm)."
            )
        if y_max > limits.y_max:
            raise ScanGeometryError(
                f"Angle {pa.angle_deg:.1f}°: scan Y range ends at {y_max:.3f} mm, "
                f"exceeds the Y axis maximum ({limits.y_max:g} mm)."
            )


def format_eta(secs: float) -> str:
    secs = max(0.0, secs)
    m, s = divmod(int(secs), 60)
    h, m = divmod(m, 60)
    if h > 0:
        return f"{h}h {m:02d}m"
    if m > 0:
        return f"{m}m {s:02d}s"
    return f"{s}s"


class EtaEstimator:
    """Rolling average of recent row durations → remaining-time estimate.

    Duration history resets when the angle index changes, since different
    angles have different row lengths.
    """

    def __init__(self, window: int = 5):
        self._durations: deque[float] = deque(maxlen=window)
        self._row_start: float | None = None
        self._last_angle_idx: int = -1

    def reset(self) -> None:
        self._durations.clear()
        self._row_start = None
        self._last_angle_idx = -1

    def row_started(self, now: float | None = None) -> None:
        self._row_start = time.monotonic() if now is None else now

    def row_finished(self, angle_idx: int, now: float | None = None) -> None:
        if angle_idx != self._last_angle_idx and self._last_angle_idx != -1:
            self._durations.clear()
        self._last_angle_idx = angle_idx
        if self._row_start is not None:
            end = time.monotonic() if now is None else now
            self._durations.append(end - self._row_start)
            self._row_start = None

    def eta_secs(self, rows_left: int) -> float | None:
        if not self._durations or rows_left <= 0:
            return None
        return sum(self._durations) / len(self._durations) * rows_left