path: root/ui/windows/plot_window.py

"""
ui/windows/plot_window.py

PLOT window — toolbar section 4.

Lets operators:
  • Choose layout: Stacked (rows) or Side-by-Side (columns) or Grid (N×M)
  • Add/remove/reorder subplot panes
  • Assign any channel (physical or derived) to any pane
  • Set X-axis: Time (elapsed s) OR any channel (e.g. strain → stress vs strain)
  • 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 typing import List, Optional, Tuple

from PyQt6.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton,
    QScrollArea, QFrame, QLineEdit, QDoubleSpinBox, QSpinBox,
    QCheckBox, QComboBox, QColorDialog, QToolButton, QButtonGroup,
    QFileDialog, QMessageBox, QRadioButton, QGroupBox, QSizePolicy,
)
from PyQt6.QtCore import Qt, pyqtSignal
from PyQt6.QtGui import QColor, QCloseEvent

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)
    layout_mode:   str            = "stacked"   # stacked | sidebyside | grid
    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)
        panes = []
        for p in d.pop("panes", []):
            traces = [TraceSpec(**t) for t in p.pop("traces", [])]
            panes.append(PaneSpec(**p, traces=traces))
        return LayoutConfig(**d, panes=panes)


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


# ══════════════════════════════════════════════════════════════════════════════
#  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()
    move_up  = pyqtSignal(object)
    move_dn  = pyqtSignal(object)

    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)

        # Header
        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))
        hl.addWidget(self._title,1)
        for txt, sig in [("▲",self.move_up),("▼",self.move_dn)]:
            b = QToolButton(); b.setText(txt); b.setObjectName("plotMoveBtn")
            b.setFixedSize(22,22); b.clicked.connect(lambda _,s=sig: s.emit(self)); hl.addWidget(b)
        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)

        # Settings row
        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:
                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"[derived] {dc.channel_id}",
                               userData=f"derived/{dc.channel_id}")
        # Select current
        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)

        # Traces
        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 Signal"); 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:
                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"[derived] {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)


# ══════════════════════════════════════════════════════════════════════════════
#  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 | Qt.WindowType.Tool)
        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, 700)
        self._blocks: List[PaneBlock] = []
        self._build()
        self._populate()

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

        # Header
        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
        glob = QWidget(); glob.setObjectName("cfgGlobalBar")
        gl   = QHBoxLayout(glob); gl.setContentsMargins(12,7,12,7); gl.setSpacing(12)

        # Layout mode
        gl.addWidget(QLabel("Layout:"))
        self._stacked_rb = QRadioButton("Stacked")
        self._side_rb    = QRadioButton("Side-by-Side")
        self._grid_rb    = QRadioButton("Grid")
        mode_grp = QButtonGroup(self)
        for rb in (self._stacked_rb, self._side_rb, self._grid_rb):
            rb.setObjectName("cfgLiveChk"); mode_grp.addButton(rb)
        mode_map = {"stacked": self._stacked_rb,
                    "sidebyside": self._side_rb, "grid": self._grid_rb}
        mode_map.get(self.cfg.layout_mode, self._stacked_rb).setChecked(True)
        for rb, key in [(self._stacked_rb,"stacked"),(self._side_rb,"sidebyside"),(self._grid_rb,"grid")]:
            rb.toggled.connect(lambda v, k=key: self._set_mode(k) if v else None)
            gl.addWidget(rb)

        gl.addWidget(QLabel("Cols:"))
        self._cols = QSpinBox(); self._cols.setRange(1,6); self._cols.setValue(self.cfg.grid_cols)
        self._cols.setFixedWidth(44); self._cols.setObjectName("traceWidthSpin")
        self._cols.valueChanged.connect(lambda v: setattr(self.cfg,"grid_cols",v))
        gl.addWidget(self._cols)

        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)

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

        # Scrollable pane blocks
        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)

        # Bottom
        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 _set_mode(self, mode: str):
        self.cfg.layout_mode = mode
        self._maybe_live()

    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)

    def _insert(self, spec: PaneSpec):
        blk = PaneBlock(spec, self.registry, self.processor)
        blk.removed.connect(self._rm); blk.changed.connect(self._maybe_live)
        blk.move_up.connect(self._mv_up); blk.move_dn.connect(self._mv_dn)
        self._blocks.append(blk)
        self._blay.insertWidget(self._blay.count()-1, blk)

    def _add_pane(self):
        spec = PaneSpec(title=f"Plot {len(self.cfg.panes)+1}")
        self.cfg.panes.append(spec); self._insert(spec); self._maybe_live()

    def _rm(self, blk):
        if blk.spec in self.cfg.panes: self.cfg.panes.remove(blk.spec)
        self._blocks.remove(blk); self._blay.removeWidget(blk); blk.deleteLater()
        self._maybe_live()

    def _mv_up(self, blk):
        i = self._blocks.index(blk)
        if i==0: return
        self.cfg.panes.insert(i-1,self.cfg.panes.pop(i))
        self._blocks.insert(i-1,self._blocks.pop(i))
        self._blay.removeWidget(blk); self._blay.insertWidget(i-1,blk)
        self._maybe_live()

    def _mv_dn(self, blk):
        i = self._blocks.index(blk)
        if i>=len(self._blocks)-1: return
        self.cfg.panes.insert(i+1,self.cfg.panes.pop(i))
        self._blocks.insert(i+1,self._blocks.pop(i))
        self._blay.removeWidget(blk); self._blay.insertWidget(i+1,blk)
        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