path: root/plugins/motion_capture/window.py

"""
motion_capture/window.py

Floating tool window for the Motion Capture plugin.

Layout:
  [Title bar]
  [FrameSelector — live / frozen feed with zoom + rubber-band target selection]
  [Status bar]
  [Controls — camera, connect, mode, shape, track/confirm, clear, px/mm]
"""

from __future__ import annotations

from typing import Optional, Tuple

import numpy as np

from PyQt6.QtCore import Qt, QPointF, QTimer, pyqtSignal
from PyQt6.QtGui import (
    QCloseEvent, QColor, QFont, QImage, QMouseEvent,
    QPainter, QPen, QPixmap,
)
from PyQt6.QtWidgets import (
    QComboBox, QDoubleSpinBox, QFrame, QHBoxLayout, QLabel, QPushButton,
    QSizePolicy, QVBoxLayout, QWidget,
)


# ── Frame display + rubber-band selection widget ──────────────────────────────

class FrameSelector(QWidget):
    """
    Live or frozen camera frame display with zoom, pan, and rubber-band
    target selection.

    Live mode  — set_frame() updates the display in real time.
    Frozen mode — freeze() locks the frame; user scrolls to zoom (around
                  cursor), right-drags to pan, left-drags to draw a
                  rubber-band selection. get_frame_selection() returns the
                  selected rect in frame pixel coordinates.
    """

    def __init__(self, parent=None):
        super().__init__(parent)
        self._frame_rgb:    Optional[np.ndarray] = None
        self._frozen_frame: Optional[np.ndarray] = None
        self._frozen   = False
        self._zoom     = 1.0
        self._pan      = QPointF(0.0, 0.0)
        self._shape    = "rect"

        self._rb_start: Optional[QPointF] = None
        self._rb_end:   Optional[QPointF] = None

        self._pan_anchor: Optional[QPointF] = None
        self._pan_origin: Optional[QPointF] = None

        self.setMouseTracking(True)
        self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
        self.setMinimumSize(480, 320)
        self.setStyleSheet("background: #0b0e13;")

    # ── Public API ────────────────────────────────────────────────────────

    def set_frame(self, rgb: np.ndarray):
        if not self._frozen:
            self._frame_rgb = rgb
            self.update()

    def set_shape(self, shape: str):
        self._shape = shape
        self.update()

    def freeze(self):
        """Lock display on current frame and enter selection mode."""
        self._frozen       = True
        self._frozen_frame = (
            self._frame_rgb.copy() if self._frame_rgb is not None else None
        )
        self._rb_start = None
        self._rb_end   = None
        self.update()

    def unfreeze(self):
        """Resume live display; reset zoom/pan/selection."""
        self._frozen       = False
        self._frozen_frame = None
        self._zoom         = 1.0
        self._pan          = QPointF(0.0, 0.0)
        self._rb_start     = None
        self._rb_end       = None
        self.update()

    def get_frozen_frame(self) -> Optional[np.ndarray]:
        return self._frozen_frame

    def get_frame_selection(self) -> Optional[Tuple[int, int, int, int]]:
        """Return (x, y, w, h) of selection in frame pixel coords, or None."""
        if self._rb_start is None or self._rb_end is None:
            return None
        frame = self._frozen_frame
        if frame is None:
            return None
        fh, fw = frame.shape[:2]
        p1 = self._screen_to_frame(self._rb_start)
        p2 = self._screen_to_frame(self._rb_end)
        x  = int(max(0,  min(p1[0], p2[0])))
        y  = int(max(0,  min(p1[1], p2[1])))
        x2 = int(min(fw, max(p1[0], p2[0])))
        y2 = int(min(fh, max(p1[1], p2[1])))
        w, h = x2 - x, y2 - y
        return (x, y, w, h) if w >= 8 and h >= 8 else None

    # ── Coordinate transforms ─────────────────────────────────────────────

    def _get_transform(self) -> Tuple[float, float, float]:
        frame = self._frozen_frame if self._frozen else self._frame_rgb
        if frame is None:
            return 1.0, 0.0, 0.0
        fh, fw = frame.shape[:2]
        W, H   = self.width(), self.height()
        ds = min(W / fw, H / fh)
        ts = ds * self._zoom
        ox = (W - fw * ts) / 2 + self._pan.x()
        oy = (H - fh * ts) / 2 + self._pan.y()
        return ts, ox, oy

    def _screen_to_frame(self, pt: QPointF) -> Tuple[float, float]:
        ts, ox, oy = self._get_transform()
        return (pt.x() - ox) / ts, (pt.y() - oy) / ts

    # ── Events ────────────────────────────────────────────────────────────

    def wheelEvent(self, e):
        frame = self._frozen_frame if self._frozen else self._frame_rgb
        if frame is None:
            e.ignore()
            return
        fh, fw = frame.shape[:2]
        W, H   = self.width(), self.height()
        ds     = min(W / fw, H / fh)

        factor   = 1.15 if e.angleDelta().y() > 0 else 1 / 1.15
        new_zoom = max(1.0, min(12.0, self._zoom * factor))
        cursor   = QPointF(e.position())

        old_ts = ds * self._zoom
        new_ts = ds * new_zoom
        old_ox = (W - fw * old_ts) / 2 + self._pan.x()
        old_oy = (H - fh * old_ts) / 2 + self._pan.y()
        fx = (cursor.x() - old_ox) / old_ts
        fy = (cursor.y() - old_oy) / old_ts
        self._pan = QPointF(
            cursor.x() - fx * new_ts - (W - fw * new_ts) / 2,
            cursor.y() - fy * new_ts - (H - fh * new_ts) / 2,
        )
        self._zoom = new_zoom
        e.accept()
        self.update()

    def mousePressEvent(self, e: QMouseEvent):
        if e.button() == Qt.MouseButton.RightButton:
            self._pan_anchor = QPointF(e.position())
            self._pan_origin = QPointF(self._pan)
        elif e.button() == Qt.MouseButton.LeftButton and self._frozen:
            self._rb_start = QPointF(e.position())
            self._rb_end   = None
            self.update()

    def mouseMoveEvent(self, e: QMouseEvent):
        if e.buttons() & Qt.MouseButton.RightButton and self._pan_anchor is not None:
            self._pan = self._pan_origin + (QPointF(e.position()) - self._pan_anchor)
            self.update()
        elif (e.buttons() & Qt.MouseButton.LeftButton
              and self._frozen and self._rb_start is not None):
            self._rb_end = QPointF(e.position())
            self.update()

    def mouseReleaseEvent(self, e: QMouseEvent):
        if e.button() == Qt.MouseButton.RightButton:
            self._pan_anchor = None
            self._pan_origin = None
        elif e.button() == Qt.MouseButton.LeftButton and self._frozen:
            if self._rb_start is not None:
                self._rb_end = QPointF(e.position())
                self.update()

    def paintEvent(self, _):
        p = QPainter(self)
        p.fillRect(self.rect(), QColor(11, 14, 19))

        frame = self._frozen_frame if self._frozen else self._frame_rgb

        if frame is None:
            p.setPen(QColor(120, 130, 150))
            p.setFont(QFont("sans-serif", 11))
            p.drawText(self.rect(), Qt.AlignmentFlag.AlignCenter,
                       "Connect a camera to begin.")
            return

        # Draw frame with zoom/pan applied
        ts, ox, oy = self._get_transform()
        fh, fw = frame.shape[:2]
        arr = np.ascontiguousarray(frame)
        qi  = QImage(arr.data, fw, fh, 3 * fw, QImage.Format.Format_RGB888)
        pix = QPixmap.fromImage(qi).scaled(
            max(1, int(fw * ts)), max(1, int(fh * ts)),
            Qt.AspectRatioMode.IgnoreAspectRatio,
            Qt.TransformationMode.FastTransformation,
        )
        p.drawPixmap(int(ox), int(oy), pix)

        # Rubber-band selection overlay
        if self._frozen and self._rb_start is not None and self._rb_end is not None:
            pen = QPen(QColor(0, 220, 100))
            pen.setWidth(2)
            pen.setStyle(Qt.PenStyle.DashLine)
            p.setPen(pen)
            sx = min(self._rb_start.x(), self._rb_end.x())
            sy = min(self._rb_start.y(), self._rb_end.y())
            rw = abs(self._rb_end.x() - self._rb_start.x())
            rh = abs(self._rb_end.y() - self._rb_start.y())
            if self._shape == "circle":
                p.drawEllipse(int(sx), int(sy), int(rw), int(rh))
            else:
                p.drawRect(int(sx), int(sy), int(rw), int(rh))

        # Instruction banner (frozen mode)
        if self._frozen:
            p.fillRect(0, 0, self.width(), 26, QColor(0, 0, 0, 170))
            p.setPen(QColor(0, 212, 255))
            p.setFont(QFont("sans-serif", 9))
            p.drawText(
                0, 0, self.width(), 26,
                Qt.AlignmentFlag.AlignCenter,
                "Scroll = zoom  ·  Right-drag = pan  ·  Left-drag = select target",
            )


# ── Main window ───────────────────────────────────────────────────────────────

class MotionCaptureWindow(QWidget):

    closed = pyqtSignal()

    def __init__(self, tracker, parent=None):
        super().__init__(parent, Qt.WindowType.Window | Qt.WindowType.Tool)
        self._tracker   = tracker
        self._connected = False
        self._selecting = False   # True while frame is frozen for selection

        self.setWindowTitle("Motion Capture")
        self.setMinimumSize(680, 560)
        self.resize(760, 600)

        self._build()

        self._refresh = QTimer(self)
        self._refresh.timeout.connect(self._update)
        self._refresh.start(33)

    # ── UI ────────────────────────────────────────────────────────────────

    def _build(self):
        root = QVBoxLayout(self)
        root.setContentsMargins(0, 0, 0, 0)
        root.setSpacing(0)

        # Title bar
        hdr = QWidget(); hdr.setObjectName("devWindowTitleBar")
        hdr.setFixedHeight(44)
        hl = QHBoxLayout(hdr); hl.setContentsMargins(14, 0, 14, 0)
        self._title_lbl = QLabel("MOTION CAPTURE")
        self._title_lbl.setObjectName("devWindowTitle")
        hl.addWidget(self._title_lbl)
        hl.addStretch()
        self._sim_badge = QLabel("SIMULATION")
        self._sim_badge.setObjectName("traceSource")
        self._sim_badge.setVisible(False)
        hl.addWidget(self._sim_badge)
        root.addWidget(hdr)

        div = QFrame(); div.setFrameShape(QFrame.Shape.HLine)
        div.setObjectName("devWindowDivider")
        root.addWidget(div)

        # Feed
        self._feed = FrameSelector()
        root.addWidget(self._feed, 1)

        # Status bar
        self._status = QLabel("Not connected")
        self._status.setObjectName("traceSource")
        self._status.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self._status.setFixedHeight(24)
        root.addWidget(self._status)

        # Controls bar
        ctrl_frame = QWidget(); ctrl_frame.setObjectName("cfgBottomBar")
        ctrl = QHBoxLayout(ctrl_frame)
        ctrl.setContentsMargins(10, 6, 10, 6); ctrl.setSpacing(6)

        # Camera selector
        ctrl.addWidget(QLabel("Camera:"))
        self._cam_combo = QComboBox()
        self._cam_combo.setObjectName("channelPickerCb")
        self._cam_combo.setMinimumWidth(180)
        ctrl.addWidget(self._cam_combo)
        self._refresh_btn = QPushButton("⟳")
        self._refresh_btn.setObjectName("configButton")
        self._refresh_btn.setFixedWidth(28)
        self._refresh_btn.setToolTip("Rescan cameras")
        self._refresh_btn.clicked.connect(self._scan_cameras)
        ctrl.addWidget(self._refresh_btn)
        self._scan_cameras()

        self._conn_btn = QPushButton("Connect")
        self._conn_btn.setObjectName("applyButton")
        self._conn_btn.clicked.connect(self._toggle_connect)
        ctrl.addWidget(self._conn_btn)

        ctrl.addWidget(_vsep())

        # Tracking mode
        ctrl.addWidget(QLabel("Mode:"))
        self._mode_cb = QComboBox()
        self._mode_cb.setObjectName("channelPickerCb")
        self._mode_cb.addItems(["Template Match", "CSRT"])
        self._mode_cb.currentIndexChanged.connect(self._on_mode_changed)
        ctrl.addWidget(self._mode_cb)

        # Shape (template mode only)
        self._shape_lbl = QLabel("Shape:")
        ctrl.addWidget(self._shape_lbl)
        self._shape_cb = QComboBox()
        self._shape_cb.setObjectName("channelPickerCb")
        self._shape_cb.addItems(["Rectangle", "Circle"])
        self._shape_cb.currentIndexChanged.connect(self._on_shape_changed)
        ctrl.addWidget(self._shape_cb)

        ctrl.addWidget(_vsep())

        # Track / confirm + clear
        self._track_btn = QPushButton("🎯  Click to Track")
        self._track_btn.setObjectName("toolbarSectionBtn")
        self._track_btn.setEnabled(False)
        self._track_btn.clicked.connect(self._on_track_btn)
        ctrl.addWidget(self._track_btn)

        self._clear_btn = QPushButton("✕  Clear")
        self._clear_btn.setObjectName("configButton")
        self._clear_btn.setEnabled(False)
        self._clear_btn.clicked.connect(self._on_clear)
        ctrl.addWidget(self._clear_btn)

        ctrl.addWidget(_vsep())

        # px/mm calibration
        ctrl.addWidget(QLabel("px/mm:"))
        self._pxmm_spin = QDoubleSpinBox()
        self._pxmm_spin.setRange(0.01, 100_000)
        self._pxmm_spin.setValue(10.0)
        self._pxmm_spin.setDecimals(2)
        self._pxmm_spin.setSuffix("  px/mm")
        self._pxmm_spin.setObjectName("cfgGlobalSpin")
        self._pxmm_spin.setFixedWidth(130)
        ctrl.addWidget(self._pxmm_spin)

        ctrl.addStretch()
        root.addWidget(ctrl_frame)

    # ── Actions ───────────────────────────────────────────────────────────

    def _scan_cameras(self):
        from tracker import list_cameras
        current_idx = self.get_camera_index()
        self._cam_combo.blockSignals(True)
        self._cam_combo.clear()
        self._cameras = list_cameras()
        for _, label in self._cameras:
            self._cam_combo.addItem(label)
        for i, (idx, _) in enumerate(self._cameras):
            if idx == current_idx:
                self._cam_combo.setCurrentIndex(i)
                break
        self._cam_combo.blockSignals(False)

    def _toggle_connect(self):
        if self._connected:
            if self._selecting:
                self._feed.unfreeze()
                self._selecting = False
            self._tracker.stop()
            self._connected = False
            self._sim_badge.setVisible(False)
            self._status.setText("Disconnected")
        else:
            ok = self._tracker.start(self.get_camera_index())
            if ok:
                self._connected = True
                self._sim_badge.setVisible(self._tracker.simulated)
            else:
                self._status.setText("Could not open camera")
        self._sync_controls()

    def _on_track_btn(self):
        if self._selecting:
            sel = self._feed.get_frame_selection()
            if sel is None:
                self._status.setText(
                    "Draw a selection box first, then click Confirm")
                return
            x, y, w, h = sel
            self._tracker.set_roi_rect(x, y, w, h,
                                        frame_rgb=self._feed.get_frozen_frame())
            self._feed.unfreeze()
            self._selecting = False
        else:
            self._feed.freeze()
            self._selecting = True
        self._sync_controls()

    def _on_clear(self):
        if self._selecting:
            self._feed.unfreeze()
            self._selecting = False
        else:
            self._tracker.clear_roi()
        self._sync_controls()

    def _on_mode_changed(self, idx: int):
        mode = "csrt" if idx == 1 else "template"
        self._tracker.set_track_mode(mode)
        is_template = mode == "template"
        self._shape_lbl.setVisible(is_template)
        self._shape_cb.setVisible(is_template)
        if self._selecting:
            self._feed.unfreeze()
            self._selecting = False
        self._sync_controls()

    def _on_shape_changed(self, idx: int):
        shape = "circle" if idx == 1 else "rect"
        self._tracker.set_template_shape(shape)
        self._feed.set_shape(shape)

    def _sync_controls(self):
        connected = self._connected
        selecting = self._selecting
        tracking  = self._tracker.is_tracking() if self._tracker else False

        self._conn_btn.setText("Disconnect" if connected else "Connect")
        self._track_btn.setEnabled(connected)
        self._clear_btn.setEnabled(connected and (selecting or tracking))

        if selecting:
            self._track_btn.setText("✓  Confirm Selection")
        else:
            self._track_btn.setText("🎯  Click to Track")

    # ── Display update ─────────────────────────────────────────────────────

    def _update(self):
        if not self._connected:
            return

        if not self._selecting:
            frame = self._tracker.get_frame()
            if frame is not None:
                self._feed.set_frame(frame)

        pos       = self._tracker.get_position()
        fps       = self._tracker.fps
        px_per_mm = self._pxmm_spin.value()

        if self._selecting:
            sel = self._feed.get_frame_selection()
            if sel:
                x, y, w, h = sel
                self._status.setText(
                    f"Selected: {w} × {h} px at ({x}, {y})  ·  Click Confirm to track")
            else:
                self._status.setText(
                    "Left-drag to select target  ·  Scroll to zoom  ·  Right-drag to pan")
        elif self._tracker.is_lost():
            self._status.setText(
                f"Tracking lost — click 'Click to Track' to reselect  |  {fps:.0f} FPS")
            self._sync_controls()
        elif pos is not None:
            xmm = pos[0] / px_per_mm
            ymm = pos[1] / px_per_mm
            self._status.setText(
                f"X: {pos[0]:.1f} px  ({xmm:.2f} mm)   "
                f"Y: {pos[1]:.1f} px  ({ymm:.2f} mm)   |   {fps:.0f} FPS")
        else:
            sim = "  [SIMULATION]" if self._tracker.simulated else ""
            self._status.setText(
                f"Live{sim}  |  {fps:.0f} FPS  |  Click 'Click to Track' to begin")

    # ── State accessors ────────────────────────────────────────────────────

    def get_camera_index(self) -> int:
        i = self._cam_combo.currentIndex()
        cameras = getattr(self, "_cameras", [])
        if 0 <= i < len(cameras):
            return cameras[i][0]
        return 0

    def set_camera_index(self, v: int):
        cameras = getattr(self, "_cameras", [])
        for i, (idx, _) in enumerate(cameras):
            if idx == v:
                self._cam_combo.setCurrentIndex(i)
                return

    def get_px_per_mm(self) -> float:
        return self._pxmm_spin.value()

    def set_px_per_mm(self, v: float):
        self._pxmm_spin.setValue(v)

    # ── Close ─────────────────────────────────────────────────────────────

    def closeEvent(self, e: QCloseEvent):
        self._refresh.stop()
        self.closed.emit()
        e.accept()


# ── Helpers ───────────────────────────────────────────────────────────────────

def _vsep() -> QFrame:
    sep = QFrame()
    sep.setFrameShape(QFrame.Shape.VLine)
    sep.setObjectName("devWindowDivider")
    return sep