sras-viewer/sras_viewer/canvases.py

"""Matplotlib canvases and the ROI primitive."""

import numpy as np
from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg
from matplotlib.figure import Figure
from matplotlib.patches import Polygon
from matplotlib.path import Path as MplPath
from PyQt6.QtCore import Qt, pyqtSignal
from PyQt6.QtGui import QKeyEvent
from PyQt6.QtWidgets import QSizePolicy

from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES, SrasFile, adc_to_mv

# ---------------------------------------------------------------------------
# ROI (free quadrilateral in data coordinates)
# ---------------------------------------------------------------------------

class RoiQuad:
    """Free quadrilateral defined in data coordinates (mm).

    Stored as 4 corner points (shape (4, 2)) in CCW order: BL, BR, TR, TL.
    Each corner can be positioned independently, allowing skewed /
    non-orthogonal regions of interest.  Because it lives in scan/data
    coords it persists unchanged when the displayed channel/mode switches.
    """

    def __init__(self, pts: np.ndarray):
        """pts : array-like, shape (4, 2)."""
        self._pts = np.asarray(pts, dtype=np.float64).reshape(4, 2).copy()

    @classmethod
    def from_bbox(cls, x0: float, y0: float, x1: float, y1: float) -> "RoiQuad":
        """Create an axis-aligned rectangle from two opposite corners."""
        lx, rx = min(x0, x1), max(x0, x1)
        by, ty = min(y0, y1), max(y0, y1)
        return cls(np.array([[lx, by], [rx, by], [rx, ty], [lx, ty]]))

    def copy(self) -> "RoiQuad":
        return RoiQuad(self._pts.copy())

    def corners(self) -> np.ndarray:
        """World-coord corners, shape (4, 2), CCW: BL, BR, TR, TL."""
        return self._pts.copy()

    def centroid(self) -> np.ndarray:
        return self._pts.mean(axis=0)

    def bbox_size(self) -> np.ndarray:
        """Width and height of the axis-aligned bounding box, shape (2,)."""
        return self._pts.max(axis=0) - self._pts.min(axis=0)

    def contains(self, x: float, y: float) -> bool:
        return bool(MplPath(self._pts).contains_point((x, y)))

    def mask_for_grid(self, x_axis: np.ndarray,
                      y_axis: np.ndarray) -> np.ndarray:
        """Boolean mask (n_rows, n_frames) of pixels whose centres lie
        inside the quadrilateral.

        Only the quad's axis-aligned bounding box is tested — meshgrid and
        contains_points over the *whole* grid would be tens of millions of
        point-in-polygon tests (and hundreds of MB of float64 temporaries)
        on a large scan, on every ROI edit.
        """
        x = np.asarray(x_axis, dtype=np.float64)
        y = np.asarray(y_axis, dtype=np.float64)
        mask = np.zeros((y.size, x.size), dtype=bool)

        (x0, y0), (x1, y1) = self._pts.min(axis=0), self._pts.max(axis=0)
        cols = np.nonzero((x >= x0) & (x <= x1))[0]
        rows = np.nonzero((y >= y0) & (y <= y1))[0]
        if cols.size == 0 or rows.size == 0:
            return mask

        c0, c1 = int(cols[0]), int(cols[-1]) + 1
        r0, r1 = int(rows[0]), int(rows[-1]) + 1
        X, Y = np.meshgrid(x[c0:c1], y[r0:r1])
        inside = MplPath(self._pts).contains_points(
            np.column_stack([X.ravel(), Y.ravel()]))
        mask[r0:r1, c0:c1] = inside.reshape(X.shape)
        return mask


# ---------------------------------------------------------------------------
# Matplotlib canvases
# ---------------------------------------------------------------------------

class ImageCanvas(FigureCanvasQTAgg):
    pixel_clicked = pyqtSignal(int, int)    # row_idx, frame_idx
    roi_changed = pyqtSignal()              # ROI created / edited / cleared
    draw_mode_changed = pyqtSignal(bool)    # "draw new ROI" arm toggled

    # Interaction state values
    _IDLE = "idle"
    _DRAW_NEW = "draw_new"
    _MOVE = "move"
    _DRAG_CORNER = "drag_corner"

    # Hit tolerance (display pixels) for handles.
    _HANDLE_PX = 12
    _CLICK_THRESH_PX = 4   # releases within this of press count as a click

    def __init__(self, parent=None):
        fig = Figure(figsize=(7, 5), tight_layout=True)
        self.ax = fig.add_subplot(111)
        super().__init__(fig)
        self.setParent(parent)
        self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
        self._extent = None
        self._img_shape = None

        # ROI state
        self._roi: RoiQuad | None = None
        self._roi_artists: list = []
        self._state = self._IDLE
        self._draw_mode = False

        # Per-interaction snapshots / anchors
        self._press_xy: tuple[float, float] | None = None
        self._press_pixel: tuple[float, float] | None = None
        self._press_button = None
        self._snapshot: RoiQuad | None = None
        self._drag_corner_idx: int = -1
        self._move_anchor = None                 # press-point in world coords
        self._draw_previous: RoiQuad | None = None

        self.mpl_connect("button_press_event", self._on_press)
        self.mpl_connect("motion_notify_event", self._on_motion)
        self.mpl_connect("button_release_event", self._on_release)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def show_image(self, img: np.ndarray, extent: list[float], cmap: str,
                   vmin: float, vmax: float, xlabel: str, ylabel: str, title: str,
                   colorbar_label: str = ""):
        self.figure.clf()
        self.ax = self.figure.add_subplot(111)
        # Patches and lines are destroyed by figure.clf(); drop stale refs.
        self._roi_artists = []

        self._extent = extent
        self._img_shape = img.shape

        im = self.ax.imshow(
            img, aspect="auto", origin="upper",
            extent=extent, cmap=cmap, vmin=vmin, vmax=vmax,
            interpolation="nearest",
        )
        cb = self.figure.colorbar(im, ax=self.ax, fraction=0.046, pad=0.04)
        if colorbar_label:
            cb.set_label(colorbar_label)

        self.ax.set_xlabel(xlabel)
        self.ax.set_ylabel(ylabel)
        self.ax.set_title(title)

        # Re-draw the ROI (if any) on top of the fresh image so it persists
        # unchanged across mode / angle / channel switches.
        self._draw_roi()
        self.draw()

    def get_roi(self) -> RoiQuad | None:
        return self._roi

    def set_roi(self, roi: RoiQuad | None):
        self._roi = roi.copy() if roi is not None else None
        self._draw_roi()
        self.draw_idle()
        self.roi_changed.emit()

    def clear_roi(self):
        self._roi = None
        self._remove_roi_artists()
        self.draw_idle()
        self.roi_changed.emit()

    def start_drawing(self):
        """Arm the next click+drag on the image to create a new ROI,
        replacing any existing one."""
        self._draw_mode = True
        self.setCursor(Qt.CursorShape.CrossCursor)
        self.draw_mode_changed.emit(True)

    def cancel_drawing(self):
        if self._draw_mode:
            self._draw_mode = False
            self.setCursor(Qt.CursorShape.ArrowCursor)
            self.draw_mode_changed.emit(False)

    # ------------------------------------------------------------------
    # Rendering
    # ------------------------------------------------------------------

    def _remove_roi_artists(self):
        for a in self._roi_artists:
            try:
                a.remove()
            except (ValueError, AttributeError, NotImplementedError):
                pass
        self._roi_artists = []

    def _draw_roi(self):
        self._remove_roi_artists()
        if self._roi is None or self.ax is None:
            return
        corners = self._roi.corners()

        # Filled quad, then a sharp unfilled edge for visibility over bright
        # images, then draggable corner handles.
        for kwargs in (
            dict(fill=True, facecolor="#ffd93a", edgecolor="#e53935",
                 alpha=0.22, linewidth=2.0, zorder=10),
            dict(fill=False, edgecolor="#e53935", linewidth=1.8, zorder=11),
        ):
            patch = Polygon(corners, closed=True, **kwargs)
            self.ax.add_patch(patch)
            self._roi_artists.append(patch)

        self._roi_artists.append(self.ax.scatter(
            corners[:, 0], corners[:, 1], s=60, c="white",
            edgecolors="#e53935", linewidths=1.6, zorder=13))

    # ------------------------------------------------------------------
    # Hit testing (display pixels for handles, data coords for "inside")
    # ------------------------------------------------------------------

    def _hit_test(self, event) -> tuple[str, int | None] | None:
        if self._roi is None or self.ax is None:
            return None
        if event.x is None or event.y is None:
            return None
        corners_disp = self.ax.transData.transform(self._roi.corners())
        click = np.array([event.x, event.y])

        for i in range(4):
            if np.hypot(*(corners_disp[i] - click)) <= self._HANDLE_PX:
                return ("corner", i)

        if event.xdata is not None and event.ydata is not None:
            if self._roi.contains(event.xdata, event.ydata):
                return ("inside", None)
        return None

    # ------------------------------------------------------------------
    # Mouse event handlers
    # ------------------------------------------------------------------

    def _on_press(self, event):
        if event.inaxes is not self.ax or self._extent is None:
            return
        if event.button != 1:   # only left mouse button
            return
        # If the matplotlib toolbar is in pan / zoom mode, let it handle
        # the interaction instead of starting a ROI manipulation.
        tb = getattr(self, "toolbar", None)
        if tb is not None and getattr(tb, "mode", ""):
            return

        self._press_xy = (event.xdata, event.ydata)
        self._press_pixel = (event.x, event.y)
        self._press_button = event.button

        if self._draw_mode:
            self._draw_previous = self._roi.copy() if self._roi else None
            self._roi = RoiQuad.from_bbox(event.xdata, event.ydata,
                                          event.xdata, event.ydata)
            self._state = self._DRAW_NEW
            self._draw_roi()
            self.draw_idle()
            return

        hit = self._hit_test(event)
        if hit is None:
            self._state = self._IDLE
            return

        kind, idx = hit
        self._snapshot = self._roi.copy()
        if kind == "corner":
            self._state = self._DRAG_CORNER
            self._drag_corner_idx = idx
        else:
            self._state = self._MOVE
            self._move_anchor = (event.xdata, event.ydata)

    def _on_motion(self, event):
        if self._state == self._IDLE:
            return
        if event.xdata is None or event.ydata is None:
            return
        if event.inaxes is not self.ax:
            return

        if self._state == self._DRAW_NEW:
            x0, y0 = self._press_xy
            self._roi = RoiQuad.from_bbox(x0, y0, event.xdata, event.ydata)
        elif self._state == self._MOVE:
            delta = np.array([event.xdata - self._move_anchor[0],
                              event.ydata - self._move_anchor[1]])
            self._roi._pts = self._snapshot.corners() + delta
        elif self._state == self._DRAG_CORNER:
            self._roi._pts[self._drag_corner_idx] = [event.xdata, event.ydata]

        self._draw_roi()
        self.draw_idle()

    def _on_release(self, event):
        if event.button != 1 and self._press_button != 1:
            return
        prev_state = self._state
        self._state = self._IDLE
        try:
            if prev_state == self._DRAW_NEW:
                self._finish_draw()
            elif prev_state in (self._MOVE, self._DRAG_CORNER):
                self._draw_roi()
                self.draw_idle()
                self.roi_changed.emit()
            else:
                self._maybe_emit_pixel_click(event)
        finally:
            self._press_xy = self._press_pixel = None
            self._press_button = None

    def _finish_draw(self):
        """Commit (or reject) a freshly-dragged quad."""
        if self._extent is not None:
            x0, x1, y_bot, y_top = self._extent
            min_w = abs(x1 - x0) * 0.01     # minimum: 1% of each axis range
            min_h = abs(y_bot - y_top) * 0.01
        else:
            min_w = min_h = 1e-6

        if self._roi is None:
            too_small = True
        else:
            bbox = self._roi.bbox_size()
            too_small = bbox[0] < min_w or bbox[1] < min_h
        if too_small:
            self._roi = self._draw_previous

        self._draw_previous = None
        self.cancel_drawing()
        self._draw_roi()
        self.draw_idle()
        self.roi_changed.emit()

    def _maybe_emit_pixel_click(self, event):
        """A release close enough to its press counts as a pixel click."""
        if (self._press_pixel is None or event.x is None or event.y is None
                or self._extent is None or event.inaxes is not self.ax
                or event.xdata is None):
            return
        dx_px = event.x - self._press_pixel[0]
        dy_px = event.y - self._press_pixel[1]
        if dx_px * dx_px + dy_px * dy_px > self._CLICK_THRESH_PX ** 2:
            return

        x0, x1, y_bot, y_top = self._extent
        n_rows, n_frames = self._img_shape
        col = int((event.xdata - x0) / (x1 - x0) * n_frames)
        row = int((event.ydata - y_top) / (y_bot - y_top) * n_rows)
        self.pixel_clicked.emit(max(0, min(row, n_rows - 1)),
                                max(0, min(col, n_frames - 1)))


class WaveformCanvas(FigureCanvasQTAgg):
    def __init__(self, parent=None):
        fig = Figure(figsize=(8, 3), tight_layout=True)
        self.ax_wave = fig.add_subplot(121)
        self.ax_right = fig.add_subplot(122)
        super().__init__(fig)
        self.setParent(parent)
        self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)

    def show_rf_waveform(self, sras: SrasFile, angle_idx: int,
                         row_idx: int, frame_idx: int,
                         apply_bg_sub: bool = True):
        """CH1 RF: time-domain + FFT spectrum.

        If apply_bg_sub is True and sras.background is not None, the background
        waveform is overlaid on the time-domain plot and the FFT is computed
        on the subtracted signal.  The unsubtracted FFT is also shown faintly
        for comparison.
        """
        data = sras.data[angle_idx]
        waveform = data[row_idx, CH1_IDX, frame_idx, :].astype(np.float32)
        t_ns = sras.time_axis_ns()
        f_mhz = sras.freq_axis_mhz()
        dc3_val = data[row_idx, CH3_IDX, frame_idx, :].astype(np.float32).mean()
        dc4_val = data[row_idx, CH4_IDX, frame_idx, :].astype(np.float32).mean()

        bg = sras.background if (apply_bg_sub and sras.background is not None) else None
        waveform_plot = waveform - bg if bg is not None else waveform

        self.ax_wave.cla()
        self.ax_right.cla()

        if bg is not None:
            self.ax_wave.plot(t_ns, waveform, linewidth=0.5, color="#aaaaaa",
                              label="raw", zorder=1)
            self.ax_wave.plot(t_ns, bg, linewidth=0.5, color="#e07030",
                              linestyle="--", label="background", zorder=2)
            self.ax_wave.plot(t_ns, waveform_plot, linewidth=0.7, color="#4488cc",
                              label="subtracted", zorder=3)
            self.ax_wave.legend(fontsize=7, loc="upper right")
        else:
            self.ax_wave.plot(t_ns, waveform, linewidth=0.7, color="#4488cc")

        self.ax_wave.set_xlabel("Time (ns)")
        self.ax_wave.set_ylabel("ADC counts")
        bg_tag = "  [bg sub]" if bg is not None else ""
        dc3_mv = adc_to_mv(dc3_val, *sras.cal(CH3_IDX))
        dc4_mv = adc_to_mv(dc4_val, *sras.cal(CH4_IDX))
        self.ax_wave.set_title(
            f"CH1 RF  row={row_idx}  frame={frame_idx}{bg_tag}\n"
            f"CH3={dc3_val:.1f}  CH4={dc4_val:.1f}  "
            f"({dc3_mv:.2f} / {dc4_mv:.2f} mV)",
            fontsize=8,
        )

        # FFT of the (possibly subtracted) waveform
        power_sub = np.abs(np.fft.rfft(waveform_plot)) ** 2
        power_sub[0] = 0.0
        peak_mhz = f_mhz[int(np.argmax(power_sub))]

        if bg is not None:
            # Also show the unsubtracted FFT for reference
            power_raw = np.abs(np.fft.rfft(waveform)) ** 2
            power_raw[0] = 0.0
            self.ax_right.plot(f_mhz, power_raw, linewidth=0.5, color="#aaaaaa",
                               label="raw FFT", zorder=1)

        self.ax_right.plot(f_mhz, power_sub, linewidth=0.7, color="#4488cc",
                           label="subtracted FFT" if bg is not None else None, zorder=2)
        self.ax_right.axvline(peak_mhz, color="tomato", linestyle="--",
                              linewidth=1.2, label=f"peak = {peak_mhz:.1f} MHz")
        self.ax_right.set_xlabel("Frequency (MHz)")
        self.ax_right.set_ylabel("Power (arb.)")
        self.ax_right.set_title("FFT Power Spectrum")
        self.ax_right.set_xlim(0, 500)
        self.ax_right.legend(fontsize=8)

        self.draw()

    def show_dc_waveform(self, sras: SrasFile, angle_idx: int, ch_idx: int,
                         row_idx: int, frame_idx: int):
        """CH3 or CH4 DC: time-domain + mean annotation."""
        waveform = sras.data[angle_idx][row_idx, ch_idx, frame_idx, :].astype(np.float32)
        mean_val = float(waveform.mean())
        mean_mv = adc_to_mv(mean_val, *sras.cal(ch_idx))

        self.ax_wave.cla()
        self.ax_right.cla()

        self.ax_wave.plot(sras.time_axis_ns(), waveform, linewidth=0.7, color="#4488cc")
        self.ax_wave.axhline(mean_val, color="tomato", linestyle="--",
                             linewidth=1.2, label=f"mean = {mean_val:.2f} ADC")
        self.ax_wave.set_xlabel("Time (ns)")
        self.ax_wave.set_ylabel("ADC counts")
        self.ax_wave.set_title(
            f"{CH_NAMES[ch_idx]} DC  row={row_idx}  frame={frame_idx}")
        self.ax_wave.legend(fontsize=8)

        self.ax_right.text(
            0.5, 0.5,
            f"DC mode\n\nmean = {mean_val:.3f} ADC\n     = {mean_mv:.3f} mV",
            ha="center", va="center",
            transform=self.ax_right.transAxes, fontsize=11,
        )
        self.ax_right.set_axis_off()

        self.draw()


class ManualAlignOverlayCanvas(FigureCanvasQTAgg):
    """Renders ManualAlignmentDialog's multi-angle mask overlay and turns
    keyboard input into translate/rotate nudge requests for whichever angle
    the dialog currently has active.

    A pure input+render widget — it holds no alignment state and never
    touches SrasFile itself; ManualAlignmentDialog owns all of that and
    decides, from these signals, whether a cheap single-layer refresh or a
    full preview-canvas rebuild is needed.

    FigureCanvasQTAgg is a real QWidget, so keyPressEvent works like on any
    other widget, but Qt only ever delivers key events to whichever widget
    currently has focus — StrongFocus, plus grabbing focus on click and once
    right after the dialog is shown, are both required or arrow keys
    silently do nothing.

    Rotate keys are letters (Q/E), not punctuation (comma/period or
    brackets): Shift+letter still reports the same Qt.Key on every platform,
    whereas Shift+comma/bracket can report a different virtual key
    (Key_Less / Key_BraceLeft) depending on platform and keyboard layout —
    which would silently break the "Shift = coarse step" modifier for
    rotation specifically. Arrow keys have no such hazard.
    """
    nudge_translate = pyqtSignal(int, int, bool)   # dir_x, dir_y in {-1,0,1}; coarse
    nudge_rotate = pyqtSignal(int, bool)           # dir in {-1,1} (CCW/CW); coarse

    _TRANSLATE_KEYS = {
        Qt.Key.Key_Left:  (-1, 0),
        Qt.Key.Key_Right: (1, 0),
        Qt.Key.Key_Up:    (0, -1),
        Qt.Key.Key_Down:  (0, 1),
    }
    _ROTATE_KEYS = {Qt.Key.Key_Q: 1, Qt.Key.Key_E: -1}   # CCW, CW

    def __init__(self, parent=None):
        fig = Figure(figsize=(6, 6), tight_layout=True)
        self.ax = fig.add_subplot(111)
        super().__init__(fig)
        self.setParent(parent)
        self.setFocusPolicy(Qt.FocusPolicy.StrongFocus)
        self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
        self.mpl_connect("button_press_event", lambda _e: self.setFocus())

    def show_overlay(self, rgba: np.ndarray, extent: list[float], title: str):
        self.figure.clf()
        self.ax = self.figure.add_subplot(111)
        self.ax.imshow(rgba, extent=extent, origin="upper", aspect="auto")
        self.ax.set_xlabel("X (mm)")
        self.ax.set_ylabel("Y (mm)")
        self.ax.set_title(title)
        self.draw_idle()   # coalesces rapid redraws — matters for key-repeat.

    def keyPressEvent(self, event: QKeyEvent):
        key = event.key()
        coarse = bool(event.modifiers() & Qt.KeyboardModifier.ShiftModifier)
        if key in self._TRANSLATE_KEYS:
            dx, dy = self._TRANSLATE_KEYS[key]
            self.nudge_translate.emit(dx, dy, coarse)
            event.accept()
        elif key in self._ROTATE_KEYS:
            self.nudge_rotate.emit(self._ROTATE_KEYS[key], coarse)
            event.accept()
        else:
            super().keyPressEvent(event)