path: root/ui/windows/plot_window.py

"""
ui/windows/plot_window.py

PLOT window — toolbar section 4.

Lets operators:
  • Arrange subplot panes by dragging tiles on the layout canvas (tiling WM style)
  • Add/remove subplot panes
  • Assign any channel (physical or derived) to any pane
  • Set X-axis: Time (elapsed s) OR any channel
  • Name axes, set Y auto/fixed range, grid, relative size weight
  • Per-trace: color, label, line style, width, visibility
  • Export/Import JSON layout

Emits layout_applied(LayoutConfig) → StripChartWidget.apply_layout()
"""

from __future__ import annotations
import json
from copy import deepcopy
from dataclasses import dataclass, field, asdict
from math import gcd
from typing import Dict, List, Optional, Tuple

from PyQt6.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton,
    QScrollArea, QFrame, QLineEdit, QDoubleSpinBox, QSpinBox,
    QCheckBox, QComboBox, QColorDialog, QToolButton,
    QFileDialog, QMessageBox,
)
from PyQt6.QtCore import Qt, pyqtSignal, QPoint, QRect
from PyQt6.QtGui import QColor, QCloseEvent, QPainter, QPen

from devices.device_registry import DeviceRegistry
from core.signal_processor import SignalProcessor


# ══════════════════════════════════════════════════════════════════════════════
#  Data model
# ══════════════════════════════════════════════════════════════════════════════

@dataclass
class TraceSpec:
    device_id:  str
    channel_id: str
    label:      str   = ""
    color:      str   = "#00d4ff"
    width:      float = 1.8
    visible:    bool  = True
    style:      str   = "solid"   # solid | dash | dot

@dataclass
class PaneSpec:
    title:    str              = "Plot"
    x_source: str              = "time"  # "time" | "dev_id/ch_id"
    y_label:  str              = ""
    x_label:  str              = "Elapsed (s)"
    y_auto:   bool             = True
    y_min:    float            = -10.0
    y_max:    float            =  10.0
    grid:     bool             = True
    weight:   int              = 1       # relative size
    traces:   List[TraceSpec]  = field(default_factory=list)

@dataclass
class LayoutConfig:
    panes:         List[PaneSpec]  = field(default_factory=list)
    tree:          Optional[dict]  = None   # BSP tree; None = legacy layout_mode
    layout_mode:   str             = "stacked"   # stacked | sidebyside | grid | free
    grid_cols:     int             = 2
    time_window_s: float           = 30.0
    link_x:        bool            = True
    show_legend:   bool            = True

    def to_json(self) -> str:
        return json.dumps(asdict(self), indent=2)

    @staticmethod
    def from_json(s: str) -> "LayoutConfig":
        d = json.loads(s)
        tree = d.pop("tree", None)
        panes = []
        for p in d.pop("panes", []):
            traces = [TraceSpec(**t) for t in p.pop("traces", [])]
            p.pop("row", None); p.pop("col", None)   # compat: old branch fields
            panes.append(PaneSpec(**p, traces=traces))
        return LayoutConfig(**d, panes=panes, tree=tree)


def build_default_layout(registry: DeviceRegistry,
                         processor: SignalProcessor) -> LayoutConfig:
    cfg = LayoutConfig()
    for dev in registry.all_instances():
        enabled = [ch for ch in dev.info.channels if ch.enabled]
        if not enabled: continue
        pane = PaneSpec(title=dev.info.name,
                        y_label=enabled[0].unit if enabled else "")
        for ch in enabled:
            pane.traces.append(TraceSpec(
                device_id=dev.info.device_id, channel_id=ch.channel_id,
                label=ch.name, color=ch.color,
            ))
        if pane.traces: cfg.panes.append(pane)
    for dc in processor.get_derived():
        pane = PaneSpec(title=dc.name, y_label=dc.unit)
        pane.traces.append(TraceSpec(
            device_id="derived", channel_id=dc.channel_id,
            label=dc.name, color=dc.color,
        ))
        cfg.panes.append(pane)
    return cfg


# ══════════════════════════════════════════════════════════════════════════════
#  BSP tree functions
# ══════════════════════════════════════════════════════════════════════════════
# Tree node schemas (as plain dicts):
#   {"kind": "leaf",   "pane": <int>}
#   {"kind": "hsplit", "ratio": <float>, "first": <node>, "second": <node>}
#   {"kind": "vsplit", "ratio": <float>, "first": <node>, "second": <node>}
# hsplit = side-by-side (first=left, second=right)
# vsplit = top-bottom   (first=top,  second=bottom)

def _tree_default_range(indices: List[int]) -> dict:
    if len(indices) == 1:
        return {"kind": "leaf", "pane": indices[0]}
    mid = len(indices) // 2
    return {"kind": "vsplit", "ratio": 0.5,
            "first":  _tree_default_range(indices[:mid]),
            "second": _tree_default_range(indices[mid:])}


def _tree_default(n: int) -> Optional[dict]:
    """Build balanced binary vsplit tree with equal ratios for n panes."""
    if n == 0:
        return None
    if n == 1:
        return {"kind": "leaf", "pane": 0}
    return _tree_equalize_ratios(_tree_default_range(list(range(n))))


def _tree_depth(node: dict) -> int:
    if node["kind"] == "leaf":
        return 0
    return 1 + max(_tree_depth(node["first"]), _tree_depth(node["second"]))


def _tree_remove(node: dict, pane_idx: int) -> Optional[dict]:
    """Remove leaf with pane_idx; promote sibling when parent loses a child."""
    if node["kind"] == "leaf":
        return None if node["pane"] == pane_idx else node
    first  = _tree_remove(node["first"],  pane_idx)
    second = _tree_remove(node["second"], pane_idx)
    if first is None and second is None:
        return None
    if first is None:
        return second
    if second is None:
        return first
    return {**node, "first": first, "second": second}


def _tree_insert(node: dict, target_idx: int, new_idx: int, zone: str) -> dict:
    """Wrap the leaf matching target_idx in a new split containing new_idx."""
    if node["kind"] == "leaf":
        if node["pane"] != target_idx:
            return node
        new_leaf = {"kind": "leaf", "pane": new_idx}
        if zone == "left":
            return {"kind": "hsplit", "ratio": 0.5, "first": new_leaf, "second": node}
        if zone == "right":
            return {"kind": "hsplit", "ratio": 0.5, "first": node, "second": new_leaf}
        if zone == "top":
            return {"kind": "vsplit", "ratio": 0.5, "first": new_leaf, "second": node}
        # bottom (default)
        return {"kind": "vsplit", "ratio": 0.5, "first": node, "second": new_leaf}
    return {**node,
            "first":  _tree_insert(node["first"],  target_idx, new_idx, zone),
            "second": _tree_insert(node["second"], target_idx, new_idx, zone)}


def _tree_reindex(node: dict, removed_idx: int) -> dict:
    """Decrement all leaf pane values > removed_idx by 1."""
    if node["kind"] == "leaf":
        pane = node["pane"]
        return {**node, "pane": pane - 1} if pane > removed_idx else node
    return {**node,
            "first":  _tree_reindex(node["first"],  removed_idx),
            "second": _tree_reindex(node["second"], removed_idx)}


def _tree_swap(node: dict, idx_a: int, idx_b: int) -> dict:
    """Swap two leaf pane values in-place (preserves tree structure)."""
    if node["kind"] == "leaf":
        if node["pane"] == idx_a: return {**node, "pane": idx_b}
        if node["pane"] == idx_b: return {**node, "pane": idx_a}
        return node
    return {**node,
            "first":  _tree_swap(node["first"],  idx_a, idx_b),
            "second": _tree_swap(node["second"], idx_a, idx_b)}


def _tree_leaf_count(node: dict) -> int:
    if node["kind"] == "leaf":
        return 1
    return _tree_leaf_count(node["first"]) + _tree_leaf_count(node["second"])


def _tree_equalize_ratios(node: dict) -> dict:
    """Reset each split's ratio so both subtrees get equal space per leaf."""
    if node["kind"] == "leaf":
        return node
    first  = _tree_equalize_ratios(node["first"])
    second = _tree_equalize_ratios(node["second"])
    n1 = _tree_leaf_count(first)
    n2 = _tree_leaf_count(second)
    return {**node, "ratio": n1 / (n1 + n2), "first": first, "second": second}


def _collect_denoms(node: dict, out: List[int]) -> None:
    if node["kind"] != "leaf":
        n1 = _tree_leaf_count(node["first"])
        n2 = _tree_leaf_count(node["second"])
        out.append(n1 + n2)
        _collect_denoms(node["first"], out)
        _collect_denoms(node["second"], out)


def tree_to_grid(node: dict, row: int, col: int, rowspan: int, colspan: int,
                 result: Dict[int, Tuple[int, int, int, int]]) -> None:
    """Recursively compute (row, col, rowspan, colspan) for each leaf."""
    if node["kind"] == "leaf":
        result[node["pane"]] = (row, col, rowspan, colspan)
    elif node["kind"] == "hsplit":
        first_cs  = max(1, round(colspan * node["ratio"]))
        second_cs = max(1, colspan - first_cs)
        first_cs  = colspan - second_cs
        tree_to_grid(node["first"],  row, col,            rowspan, first_cs,  result)
        tree_to_grid(node["second"], row, col + first_cs, rowspan, second_cs, result)
    elif node["kind"] == "vsplit":
        first_rs  = max(1, round(rowspan * node["ratio"]))
        second_rs = max(1, rowspan - first_rs)
        first_rs  = rowspan - second_rs
        tree_to_grid(node["first"],  row,            col, first_rs,  colspan, result)
        tree_to_grid(node["second"], row + first_rs, col, second_rs, colspan, result)


def _tree_grid_size(node: dict) -> int:
    """Minimal grid = LCM of all split denominators (leaf counts), capped at 12."""
    denoms: List[int] = []
    _collect_denoms(node, denoms)
    if not denoms:
        return 1
    result = 1
    for d in denoms:
        result = result * d // gcd(result, d)
        if result >= 12:
            return 12
    return result


# ══════════════════════════════════════════════════════════════════════════════
#  Color swatch
# ══════════════════════════════════════════════════════════════════════════════

class ColorSwatch(QPushButton):
    color_changed = pyqtSignal(str)
    def __init__(self, color="#00d4ff"):
        super().__init__(); self._c = color
        self.setFixedSize(26,20); self._paint(); self.clicked.connect(self._pick)
    def _paint(self):
        self.setStyleSheet(f"QPushButton{{background:{self._c};border:1px solid #2a3558;"
                           f"border-radius:3px;}}QPushButton:hover{{border-color:#3b82f6;}}")
    def _pick(self):
        c = QColorDialog.getColor(QColor(self._c), self)
        if c.isValid(): self._c = c.name(); self._paint(); self.color_changed.emit(self._c)
    @property
    def color(self): return self._c


# ══════════════════════════════════════════════════════════════════════════════
#  Trace row
# ══════════════════════════════════════════════════════════════════════════════

class TraceRow(QFrame):
    removed = pyqtSignal(object)
    changed = pyqtSignal()
    def __init__(self, tr: TraceSpec):
        super().__init__(); self.tr = tr; self.setObjectName("traceRow"); self._build()
    def _build(self):
        lay = QHBoxLayout(self); lay.setContentsMargins(6,3,6,3); lay.setSpacing(6)
        self._vis = QCheckBox(); self._vis.setChecked(self.tr.visible)
        self._vis.toggled.connect(lambda v: (setattr(self.tr,"visible",v), self.changed.emit()))
        lay.addWidget(self._vis)
        self._sw = ColorSwatch(self.tr.color)
        self._sw.color_changed.connect(lambda c: (setattr(self.tr,"color",c), self.changed.emit()))
        lay.addWidget(self._sw)
        src = QLabel(f"{self.tr.device_id} / {self.tr.channel_id}")
        src.setObjectName("traceSource"); src.setMinimumWidth(130); lay.addWidget(src)
        self._lbl = QLineEdit(self.tr.label); self._lbl.setObjectName("traceLabel")
        self._lbl.setPlaceholderText("Label…")
        self._lbl.textChanged.connect(lambda t: setattr(self.tr,"label",t))
        lay.addWidget(self._lbl,1)
        self._sty = QComboBox(); self._sty.setObjectName("traceStyleCb")
        self._sty.addItems(["solid","dash","dot"]); self._sty.setCurrentText(self.tr.style)
        self._sty.setFixedWidth(64)
        self._sty.currentTextChanged.connect(lambda s: setattr(self.tr,"style",s))
        lay.addWidget(self._sty)
        self._w = QDoubleSpinBox(); self._w.setObjectName("traceWidthSpin")
        self._w.setRange(0.5,6.0); self._w.setSingleStep(0.5)
        self._w.setValue(self.tr.width); self._w.setFixedWidth(56)
        self._w.valueChanged.connect(lambda v: setattr(self.tr,"width",v))
        lay.addWidget(self._w)
        rm = QToolButton(); rm.setText("✕"); rm.setObjectName("traceRemoveBtn")
        rm.setFixedSize(22,22); rm.clicked.connect(lambda: self.removed.emit(self))
        lay.addWidget(rm)


# ══════════════════════════════════════════════════════════════════════════════
#  Pane block
# ══════════════════════════════════════════════════════════════════════════════

class PaneBlock(QFrame):
    removed       = pyqtSignal(object)
    changed       = pyqtSignal()
    title_changed = pyqtSignal(object, str)   # (self, new_title)

    def __init__(self, spec: PaneSpec, registry: DeviceRegistry,
                 processor: SignalProcessor):
        super().__init__()
        self.spec = spec; self.registry = registry; self.processor = processor
        self._rows: List[TraceRow] = []
        self.setObjectName("plotBlock"); self._build()

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

        hdr = QWidget(); hdr.setObjectName("plotBlockHeader"); hdr.setFixedHeight(32)
        hl = QHBoxLayout(hdr); hl.setContentsMargins(8,0,6,0); hl.setSpacing(4)
        self._title = QLineEdit(self.spec.title); self._title.setObjectName("plotBlockTitle")
        self._title.textChanged.connect(lambda t: (
            setattr(self.spec, "title", t),
            self.title_changed.emit(self, t),
        ))
        hl.addWidget(self._title, 1)
        rm = QToolButton(); rm.setText("✕"); rm.setObjectName("plotRemoveBtn")
        rm.setFixedSize(22,22); rm.clicked.connect(lambda: self.removed.emit(self))
        hl.addWidget(rm)
        outer.addWidget(hdr)

        sett = QWidget(); sett.setObjectName("plotBlockSettings")
        sl = QHBoxLayout(sett); sl.setContentsMargins(8,6,8,6); sl.setSpacing(8)

        sl.addWidget(QLabel("X Axis:"))
        self._x_cb = QComboBox(); self._x_cb.setObjectName("channelPickerCb")
        self._x_cb.setMinimumWidth(150)
        self._x_cb.addItem("⏱  Time (elapsed s)", userData="time")
        for dev in self.registry.all_instances():
            for ch in dev.info.channels:
                if not ch.enabled: continue
                self._x_cb.addItem(f"{dev.info.device_id}/{ch.channel_id}  ({ch.name})",
                                   userData=f"{dev.info.device_id}/{ch.channel_id}")
        for dc in self.processor.get_derived():
            self._x_cb.addItem(f"[virtual] {dc.channel_id}",
                               userData=f"derived/{dc.channel_id}")
        for i in range(self._x_cb.count()):
            if self._x_cb.itemData(i) == self.spec.x_source:
                self._x_cb.setCurrentIndex(i); break
        self._x_cb.currentIndexChanged.connect(
            lambda: setattr(self.spec,"x_source", self._x_cb.currentData() or "time"))
        sl.addWidget(self._x_cb)

        sl.addWidget(QLabel("Y:"))
        self._ylabel = QLineEdit(self.spec.y_label); self._ylabel.setObjectName("plotYLabelEdit")
        self._ylabel.setPlaceholderText("Y label"); self._ylabel.setMaximumWidth(110)
        self._ylabel.textChanged.connect(lambda t: setattr(self.spec,"y_label",t))
        sl.addWidget(self._ylabel)

        self._auto = QCheckBox("Auto Y"); self._auto.setChecked(self.spec.y_auto)
        self._auto.toggled.connect(self._on_auto); sl.addWidget(self._auto)

        sl.addWidget(QLabel("Min:"))
        self._ymin = QDoubleSpinBox(); self._ymin.setRange(-1e9,1e9)
        self._ymin.setValue(self.spec.y_min); self._ymin.setFixedWidth(72)
        self._ymin.setEnabled(not self.spec.y_auto)
        self._ymin.valueChanged.connect(lambda v: setattr(self.spec,"y_min",v))
        sl.addWidget(self._ymin)

        sl.addWidget(QLabel("Max:"))
        self._ymax = QDoubleSpinBox(); self._ymax.setRange(-1e9,1e9)
        self._ymax.setValue(self.spec.y_max); self._ymax.setFixedWidth(72)
        self._ymax.setEnabled(not self.spec.y_auto)
        self._ymax.valueChanged.connect(lambda v: setattr(self.spec,"y_max",v))
        sl.addWidget(self._ymax)

        self._grid = QCheckBox("Grid"); self._grid.setChecked(self.spec.grid)
        self._grid.toggled.connect(lambda v: setattr(self.spec,"grid",v))
        sl.addWidget(self._grid)

        sl.addWidget(QLabel("Size:"))
        self._wt = QSpinBox(); self._wt.setRange(1,10); self._wt.setValue(self.spec.weight)
        self._wt.setFixedWidth(44); self._wt.setToolTip("Relative pane size")
        self._wt.valueChanged.connect(lambda v: setattr(self.spec,"weight",v))
        sl.addWidget(self._wt); sl.addStretch()
        outer.addWidget(sett)

        tr_w = QWidget(); tr_w.setObjectName("plotBlockTraces")
        self._tlay = QVBoxLayout(tr_w)
        self._tlay.setContentsMargins(4,2,4,4); self._tlay.setSpacing(2)
        for tr in self.spec.traces: self._add_row(tr)

        pick_row = QHBoxLayout()
        self._picker = self._make_picker(); pick_row.addWidget(self._picker,1)
        ab = QPushButton("＋ Add Channel"); ab.setObjectName("addTraceBtn")
        ab.clicked.connect(self._on_add); pick_row.addWidget(ab)
        self._tlay.addLayout(pick_row)
        outer.addWidget(tr_w)

    def _make_picker(self):
        cb = QComboBox(); cb.setObjectName("channelPickerCb")
        for dev in self.registry.all_instances():
            for ch in dev.info.channels:
                if not ch.enabled: continue
                cb.addItem(f"{dev.info.device_id} / {ch.channel_id}  ({ch.name})",
                           userData=(dev.info.device_id, ch.channel_id, ch.name, ch.color))
        for dc in self.processor.get_derived():
            cb.addItem(f"[virtual] {dc.channel_id}  ({dc.name})",
                       userData=("derived", dc.channel_id, dc.name, dc.color))
        return cb

    def _add_row(self, tr: TraceSpec):
        row = TraceRow(tr); row.removed.connect(self._rm_row)
        row.changed.connect(self.changed)
        self._rows.append(row)
        self._tlay.insertWidget(self._tlay.count()-1, row)

    def _rm_row(self, row: TraceRow):
        if row.tr in self.spec.traces: self.spec.traces.remove(row.tr)
        self._rows.remove(row); self._tlay.removeWidget(row); row.deleteLater()
        self.changed.emit()

    def _on_add(self):
        data = self._picker.currentData()
        if not data: return
        d,c,n,col = data
        tr = TraceSpec(device_id=d,channel_id=c,label=n,color=col)
        self.spec.traces.append(tr); self._add_row(tr); self.changed.emit()

    def _on_auto(self, v):
        self.spec.y_auto = v; self._ymin.setEnabled(not v); self._ymax.setEnabled(not v)


# ══════════════════════════════════════════════════════════════════════════════
#  Layout canvas — drag-to-split BSP tiling
# ══════════════════════════════════════════════════════════════════════════════

_LC_GAP      = 5    # gap between tiles in pixels
_LC_MIN_TILE = 40   # minimum tile dimension in pixels
_LC_HEIGHT   = 160  # fixed canvas height


def _zone_overlay_rect(r: QRect, zone: str) -> QRect:
    if zone == "left":
        return QRect(r.x(), r.y(), r.width() // 2, r.height())
    if zone == "right":
        return QRect(r.x() + r.width() // 2, r.y(), r.width() - r.width() // 2, r.height())
    if zone == "top":
        return QRect(r.x(), r.y(), r.width(), r.height() // 2)
    if zone == "bottom":
        return QRect(r.x(), r.y() + r.height() // 2, r.width(), r.height() - r.height() // 2)
    # center (move)
    return r


class PaneTile(QFrame):
    def __init__(self, pane_index: int, title: str, parent: QWidget) -> None:
        super().__init__(parent)
        self.pane_index = pane_index
        self.setObjectName("layoutTile")
        self.setAttribute(Qt.WidgetAttribute.WA_TransparentForMouseEvents, True)
        lay = QVBoxLayout(self)
        lay.setContentsMargins(4, 4, 4, 4)
        self._lbl = QLabel(title)
        self._lbl.setObjectName("layoutTileLabel")
        self._lbl.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self._lbl.setWordWrap(True)
        lay.addWidget(self._lbl)

    def set_title(self, title: str) -> None:
        self._lbl.setText(title)

    def set_dragging(self, v: bool) -> None:
        name = "layoutTileDragging" if v else "layoutTile"
        self.setObjectName(name)
        self.style().unpolish(self); self.style().polish(self)
        self._lbl.style().unpolish(self._lbl); self._lbl.style().polish(self._lbl)


class LayoutCanvas(QFrame):
    arrangement_changed = pyqtSignal()

    def __init__(self, panes: List[PaneSpec], tree: Optional[dict],
                 parent: Optional[QWidget] = None) -> None:
        super().__init__(parent)
        self.setObjectName("layoutCanvasWidget")
        self.setMouseTracking(True)
        self.setFixedHeight(_LC_HEIGHT)
        self._cfg_ref: Optional[LayoutConfig] = None   # set by PlotWindow
        self._panes: List[PaneSpec] = panes
        self._tree: Optional[dict] = None
        self._tiles: List[PaneTile] = []
        self._tile_rects: Dict[int, QRect] = {}
        self._drag_tile:   Optional[PaneTile]       = None
        self._drag_offset: QPoint                   = QPoint()
        self._hover_zone:  Optional[Tuple[int, str]] = None
        self.sync_panes(panes, tree)

    # ── Public ────────────────────────────────────────────────────────────────

    def sync_panes(self, panes: List[PaneSpec], tree: Optional[dict]) -> None:
        self._panes = panes
        self._drag_tile = None
        self._hover_zone = None
        n = len(panes)
        if tree is None or not _tree_valid(tree, n):
            tree = _tree_default(n)
        self._tree = tree
        self._rebuild_tiles()

    def update_tile_title(self, pane_index: int, title: str) -> None:
        if 0 <= pane_index < len(self._tiles):
            self._tiles[pane_index].set_title(title)

    # ── Geometry ──────────────────────────────────────────────────────────────

    def resizeEvent(self, event) -> None:
        super().resizeEvent(event)
        self._reposition_tiles()

    def _rebuild_tiles(self) -> None:
        for t in self._tiles:
            t.deleteLater()
        self._tiles.clear()
        for i, pane in enumerate(self._panes):
            tile = PaneTile(i, pane.title, self)
            tile.show()
            self._tiles.append(tile)
        self._reposition_tiles()

    def _reposition_tiles(self) -> None:
        if not self._tree or not self._tiles:
            self._tile_rects = {}
            return
        canvas = QRect(_LC_GAP, _LC_GAP,
                       self.width() - 2 * _LC_GAP,
                       self.height() - 2 * _LC_GAP)
        self._tile_rects = {}
        self._walk_rects(self._tree, canvas)
        for tile in self._tiles:
            r = self._tile_rects.get(tile.pane_index)
            if r:
                tile.setGeometry(r)

    def _walk_rects(self, node: dict, rect: QRect) -> None:
        if node["kind"] == "leaf":
            self._tile_rects[node["pane"]] = rect
        elif node["kind"] == "hsplit":
            left_w = max(_LC_MIN_TILE, int(rect.width() * node["ratio"]))
            right_w = rect.width() - left_w - _LC_GAP
            right_w = max(_LC_MIN_TILE, right_w)
            left_w = rect.width() - right_w - _LC_GAP
            self._walk_rects(node["first"],
                             QRect(rect.x(), rect.y(), left_w, rect.height()))
            self._walk_rects(node["second"],
                             QRect(rect.x() + left_w + _LC_GAP, rect.y(), right_w, rect.height()))
        elif node["kind"] == "vsplit":
            top_h = max(_LC_MIN_TILE, int(rect.height() * node["ratio"]))
            bot_h = rect.height() - top_h - _LC_GAP
            bot_h = max(_LC_MIN_TILE, bot_h)
            top_h = rect.height() - bot_h - _LC_GAP
            self._walk_rects(node["first"],
                             QRect(rect.x(), rect.y(), rect.width(), top_h))
            self._walk_rects(node["second"],
                             QRect(rect.x(), rect.y() + top_h + _LC_GAP, rect.width(), bot_h))

    # ── Drop zone detection ───────────────────────────────────────────────────

    def _zone_at(self, pos: QPoint) -> Optional[Tuple[int, str]]:
        dragging_idx = self._drag_tile.pane_index if self._drag_tile else -1
        for pane_idx, rect in self._tile_rects.items():
            if pane_idx == dragging_idx:
                continue
            if not rect.contains(pos):
                continue
            rel_x = pos.x() - rect.x()
            rel_y = pos.y() - rect.y()
            w, h = rect.width(), rect.height()
            zone_w = max(1, w // 3)
            zone_h = max(1, h // 3)
            if rel_x < zone_w:
                return (pane_idx, "left")
            if rel_x > w - zone_w:
                return (pane_idx, "right")
            if rel_y < zone_h:
                return (pane_idx, "top")
            if rel_y > h - zone_h:
                return (pane_idx, "bottom")
            return (pane_idx, "center")  # center — move/swap
        return None

    # ── Drag ─────────────────────────────────────────────────────────────────

    def mousePressEvent(self, event) -> None:
        if event.button() != Qt.MouseButton.LeftButton:
            return
        for tile in reversed(self._tiles):
            if tile.geometry().contains(event.pos()):
                self._drag_tile = tile
                self._drag_offset = event.pos() - tile.pos()
                tile.set_dragging(True)
                tile.raise_()
                return

    def mouseMoveEvent(self, event) -> None:
        if self._drag_tile is None:
            return
        self._drag_tile.move(event.pos() - self._drag_offset)
        zone = self._zone_at(event.pos())
        if zone != self._hover_zone:
            self._hover_zone = zone
            self.update()

    def mouseReleaseEvent(self, event) -> None:
        if self._drag_tile is None:
            return
        dragged_idx = self._drag_tile.pane_index
        zone_info = self._zone_at(event.pos())
        self._drag_tile.set_dragging(False)
        self._drag_tile = None
        self._hover_zone = None

        if zone_info is not None:
            target_idx, zone = zone_info
            self._do_drop(dragged_idx, target_idx, zone)
        else:
            # No valid drop target — restore tile positions
            self._reposition_tiles()
        self.update()

    def _do_drop(self, dragged_idx: int, target_idx: int, zone: str) -> None:
        if dragged_idx == target_idx or self._tree is None:
            self._reposition_tiles()
            return
        if zone == "center":
            new_tree = _tree_swap(self._tree, dragged_idx, target_idx)
        else:
            new_tree = _tree_remove(self._tree, dragged_idx)
            if new_tree is None:
                self._reposition_tiles()
                return
            new_tree = _tree_equalize_ratios(_tree_insert(new_tree, target_idx, dragged_idx, zone))
        self._tree = new_tree
        if self._cfg_ref is not None:
            self._cfg_ref.tree = new_tree
        self._reposition_tiles()
        self.arrangement_changed.emit()

    # ── Paint ─────────────────────────────────────────────────────────────────

    def paintEvent(self, event) -> None:
        super().paintEvent(event)
        if self._hover_zone is None or self._drag_tile is None:
            return
        pane_idx, zone = self._hover_zone
        r = self._tile_rects.get(pane_idx)
        if r is None:
            return
        overlay = _zone_overlay_rect(r, zone)
        painter = QPainter(self)
        painter.setRenderHint(QPainter.RenderHint.Antialiasing)
        if zone == "center":
            painter.setPen(QPen(QColor("#22c55e"), 2))
            painter.setBrush(QColor(34, 197, 94, 60))
        else:
            painter.setPen(QPen(QColor("#3b82f6"), 2))
            painter.setBrush(QColor(59, 130, 246, 70))
        painter.drawRoundedRect(overlay, 4, 4)


def _tree_valid(tree: dict, n_panes: int) -> bool:
    """Check tree references only valid pane indices."""
    try:
        leaves = []
        _collect_leaves(tree, leaves)
        return sorted(leaves) == list(range(n_panes))
    except Exception:
        return False


def _collect_leaves(node: dict, out: list) -> None:
    if node["kind"] == "leaf":
        out.append(node["pane"])
    else:
        _collect_leaves(node["first"], out)
        _collect_leaves(node["second"], out)


# ══════════════════════════════════════════════════════════════════════════════
#  PlotWindow
# ══════════════════════════════════════════════════════════════════════════════

class PlotWindow(QWidget):
    layout_applied = pyqtSignal(object)   # LayoutConfig
    closed         = pyqtSignal()

    def __init__(self, registry: DeviceRegistry, processor: SignalProcessor,
                 current: Optional[LayoutConfig] = None, parent=None):
        super().__init__(parent, Qt.WindowType.Window)
        self.registry  = registry
        self.processor = processor
        self.cfg       = deepcopy(current) if current \
                         else build_default_layout(registry, processor)
        self.setWindowTitle("Plot Builder")
        self.setMinimumSize(720, 560)
        self.resize(820, 760)
        self._blocks: List[PaneBlock] = []
        self._build()
        self._populate()

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

        hdr = QWidget(); hdr.setObjectName("devWindowTitleBar"); hdr.setFixedHeight(44)
        hl  = QHBoxLayout(hdr); hl.setContentsMargins(14,0,10,0)
        hl.addWidget(QLabel("PLOT BUILDER").also(lambda w: w.setObjectName("devWindowTitle")),1)
        root.addWidget(hdr)
        div = QFrame(); div.setFrameShape(QFrame.Shape.HLine)
        div.setObjectName("devWindowDivider"); root.addWidget(div)

        # Global settings bar
        glob = QWidget(); glob.setObjectName("cfgGlobalBar")
        gl   = QHBoxLayout(glob); gl.setContentsMargins(12,7,12,7); gl.setSpacing(12)
        gl.addWidget(QLabel("Window:"))
        self._win = QDoubleSpinBox(); self._win.setRange(1,3600)
        self._win.setSuffix(" s"); self._win.setValue(self.cfg.time_window_s)
        self._win.setObjectName("cfgGlobalSpin")
        self._win.valueChanged.connect(lambda v: setattr(self.cfg,"time_window_s",v))
        gl.addWidget(self._win)
        self._linkx = QCheckBox("Link X"); self._linkx.setChecked(self.cfg.link_x)
        self._linkx.setObjectName("cfgLiveChk")
        self._linkx.toggled.connect(lambda v: setattr(self.cfg,"link_x",v))
        gl.addWidget(self._linkx)
        self._legend = QCheckBox("Legend"); self._legend.setChecked(self.cfg.show_legend)
        self._legend.setObjectName("cfgLiveChk")
        self._legend.toggled.connect(lambda v: setattr(self.cfg,"show_legend",v))
        gl.addWidget(self._legend)
        self._live = QCheckBox("Live preview"); self._live.setObjectName("cfgLiveChk")
        gl.addWidget(self._live)
        gl.addStretch()
        add_btn = QPushButton("＋  Add Pane"); add_btn.setObjectName("addDeviceButton")
        add_btn.clicked.connect(self._add_pane); gl.addWidget(add_btn)
        root.addWidget(glob)

        # Layout canvas section
        canvas_bar = QWidget(); canvas_bar.setObjectName("layoutCanvasBar")
        cb_lay = QVBoxLayout(canvas_bar)
        cb_lay.setContentsMargins(10, 6, 10, 8); cb_lay.setSpacing(4)
        clbl = QLabel("LAYOUT  —  drag tiles to edges to split, center to move")
        clbl.setObjectName("layoutCanvasLabel")
        cb_lay.addWidget(clbl)
        self._canvas = LayoutCanvas(self.cfg.panes, self.cfg.tree, self)
        self._canvas._cfg_ref = self.cfg
        self._canvas.arrangement_changed.connect(self._maybe_live)
        cb_lay.addWidget(self._canvas)
        root.addWidget(canvas_bar)

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

        scroll = QScrollArea(); scroll.setWidgetResizable(True)
        scroll.setObjectName("deviceScroll")
        self._cont = QWidget()
        self._blay = QVBoxLayout(self._cont)
        self._blay.setContentsMargins(10,10,10,10); self._blay.setSpacing(12)
        self._blay.addStretch()
        scroll.setWidget(self._cont)
        root.addWidget(scroll, 1)

        btm = QWidget(); btm.setObjectName("cfgBottomBar")
        bl  = QHBoxLayout(btm); bl.setContentsMargins(12,8,12,8)
        for lbl, fn in [("⟳  Rebuild from Channels", self._reset),
                         ("↓  Export JSON", self._export),
                         ("↑  Import JSON", self._import_json)]:
            b = QPushButton(lbl); b.setObjectName("configButton")
            b.clicked.connect(fn); bl.addWidget(b)
        bl.addStretch()
        ap = QPushButton("✓  Apply"); ap.setObjectName("applyButton")
        ap.clicked.connect(self._apply); bl.addWidget(ap)
        root.addWidget(btm)

    def _populate(self):
        for b in self._blocks:
            self._blay.removeWidget(b); b.deleteLater()
        self._blocks.clear()
        for spec in self.cfg.panes:
            self._insert(spec)
        self.cfg.layout_mode = "free"
        self._canvas._cfg_ref = self.cfg
        self._canvas.sync_panes(self.cfg.panes, self.cfg.tree)

    def _insert(self, spec: PaneSpec):
        blk = PaneBlock(spec, self.registry, self.processor)
        blk.removed.connect(self._rm)
        blk.changed.connect(self._maybe_live)
        blk.title_changed.connect(self._on_tile_title)
        self._blocks.append(blk)
        self._blay.insertWidget(self._blay.count()-1, blk)

    def _on_tile_title(self, blk: PaneBlock, title: str) -> None:
        if blk in self._blocks:
            self._canvas.update_tile_title(self._blocks.index(blk), title)

    def _add_pane(self):
        spec = PaneSpec(title=f"Plot {len(self.cfg.panes)+1}")
        self.cfg.panes.append(spec)
        new_idx = len(self.cfg.panes) - 1
        if self.cfg.tree is None:
            self.cfg.tree = {"kind": "leaf", "pane": 0}
        else:
            self.cfg.tree = _tree_equalize_ratios({
                "kind": "vsplit", "ratio": 0.5,
                "first": self.cfg.tree,
                "second": {"kind": "leaf", "pane": new_idx},
            })
        self._insert(spec)
        self.cfg.layout_mode = "free"
        self._canvas.sync_panes(self.cfg.panes, self.cfg.tree)
        self._maybe_live()

    def _rm(self, blk: PaneBlock):
        removed_idx = self._blocks.index(blk)
        self.cfg.panes.remove(blk.spec)
        self._blocks.remove(blk)
        self._blay.removeWidget(blk); blk.deleteLater()
        if self.cfg.tree:
            self.cfg.tree = _tree_remove(self.cfg.tree, removed_idx)
            if self.cfg.tree:
                self.cfg.tree = _tree_equalize_ratios(
                    _tree_reindex(self.cfg.tree, removed_idx)
                )
        self._canvas.sync_panes(self.cfg.panes, self.cfg.tree)
        self._maybe_live()

    def _maybe_live(self):
        if self._live.isChecked():
            self.layout_applied.emit(deepcopy(self.cfg))

    def _apply(self):
        self.layout_applied.emit(deepcopy(self.cfg))

    def _reset(self):
        self.cfg = build_default_layout(self.registry, self.processor)
        self._win.setValue(self.cfg.time_window_s)
        self._populate()
        self._maybe_live()

    def _export(self):
        path, _ = QFileDialog.getSaveFileName(self, "Export", "layout.json", "JSON (*.json)")
        if path: open(path,"w").write(self.cfg.to_json())

    def _import_json(self):
        path, _ = QFileDialog.getOpenFileName(self, "Import", "", "JSON (*.json)")
        if path:
            try:
                self.cfg = LayoutConfig.from_json(open(path).read())
                self._win.setValue(self.cfg.time_window_s)
                self._populate()
                self._maybe_live()
            except Exception as e:
                QMessageBox.critical(self, "Import failed", str(e))

    def refresh_channels(self):
        self._populate()

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


# monkey-patch QLabel.also for inline objectName setting
from PyQt6.QtWidgets import QLabel as _QLabel2
def _also2(self, fn): fn(self); return self
if not hasattr(_QLabel2, "also"): _QLabel2.also = _also2