path: root/ui/plot_config.py

"""
ui/plot_config.py

Plot Configuration Builder — floating window.

Lets the user define:
  • How many subplot rows to show
  • Which channels go in each row (any mix of devices/channels)
  • Per-channel: label override, color, line width, visibility
  • Per-plot:    Y-axis label, Y min/max (or auto), grid on/off
  • Global:      time window, sample decimation, background color

The resulting config is stored as a list of PlotConfig dataclasses
and emitted via the `config_applied` signal so StripChartWidget can
rebuild itself on demand.
"""

from __future__ import annotations
import json
from copy import deepcopy
from dataclasses import dataclass, field, asdict
from typing import List, Optional

from PyQt6.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton,
    QScrollArea, QFrame, QLineEdit, QDoubleSpinBox, QSpinBox,
    QCheckBox, QComboBox, QColorDialog, QGroupBox, QFormLayout,
    QSizePolicy, QSplitter, QListWidget, QListWidgetItem,
    QAbstractItemView, QToolButton, QSlider,
)
from PyQt6.QtCore import Qt, pyqtSignal, QSize
from PyQt6.QtGui import QColor, QIcon, QCloseEvent

from devices.device_registry import DeviceRegistry


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

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


@dataclass
class PlotConfig:
    title:      str                      = "Plot"
    y_label:    str                      = ""
    y_auto:     bool                     = True
    y_min:      float                    = -10.0
    y_max:      float                    =  10.0
    grid:       bool                     = True
    traces:     List[ChannelTrace]       = field(default_factory=list)


@dataclass
class ChartConfig:
    plots:          List[PlotConfig]     = field(default_factory=list)
    time_window_s:  float                = 30.0
    bg_color:       str                  = "#0b0e13"
    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) -> "ChartConfig":
        d = json.loads(s)
        plots = []
        for p in d.get("plots", []):
            traces = [ChannelTrace(**t) for t in p.pop("traces", [])]
            plots.append(PlotConfig(**p, traces=traces))
        d["plots"] = plots
        return ChartConfig(**{k: v for k, v in d.items() if k != "plots"},
                           plots=plots)


def default_config(registry: DeviceRegistry) -> ChartConfig:
    """Build a sensible default config from whatever devices are registered."""
    cfg = ChartConfig()
    for dev in registry.all_instances():
        if not dev.info.channels:
            continue
        p = PlotConfig(
            title=dev.info.name,
            y_label=dev.info.channels[0].unit if dev.info.channels else "",
        )
        for ch in dev.info.channels:
            if ch.enabled:
                p.traces.append(ChannelTrace(
                    device_id=dev.info.device_id,
                    channel_id=ch.channel_id,
                    label=ch.name,
                    color=ch.color,
                ))
        if p.traces:
            cfg.plots.append(p)
    return cfg


# ══════════════════════════════════════════════════════════════════════════════
#  Colour swatch button
# ══════════════════════════════════════════════════════════════════════════════

class ColorButton(QPushButton):
    color_changed = pyqtSignal(str)

    def __init__(self, color: str = "#00d4ff"):
        super().__init__()
        self._color = color
        self.setFixedSize(28, 22)
        self._apply()
        self.clicked.connect(self._pick)

    def _apply(self):
        self.setStyleSheet(
            f"QPushButton {{ background:{self._color}; border:1px solid #2a3558;"
            f" border-radius:3px; }} "
            f"QPushButton:hover {{ border-color:#3b82f6; }}"
        )

    def _pick(self):
        c = QColorDialog.getColor(QColor(self._color), self, "Pick colour")
        if c.isValid():
            self._color = c.name()
            self._apply()
            self.color_changed.emit(self._color)

    @property
    def color(self): return self._color

    @color.setter
    def color(self, v: str):
        self._color = v
        self._apply()


# ══════════════════════════════════════════════════════════════════════════════
#  Trace row widget  (one row per ChannelTrace inside a plot)
# ══════════════════════════════════════════════════════════════════════════════

class TraceRow(QFrame):
    removed   = pyqtSignal(object)   # self
    changed   = pyqtSignal()

    def __init__(self, trace: ChannelTrace, registry: DeviceRegistry):
        super().__init__()
        self.trace    = trace
        self.registry = registry
        self.setObjectName("traceRow")
        self._build()

    def _build(self):
        lay = QHBoxLayout(self)
        lay.setContentsMargins(6, 4, 6, 4)
        lay.setSpacing(6)

        # Visibility checkbox
        self._vis = QCheckBox()
        self._vis.setChecked(self.trace.visible)
        self._vis.setToolTip("Show/hide trace")
        self._vis.toggled.connect(self._on_vis)
        lay.addWidget(self._vis)

        # Color swatch
        self._color_btn = ColorButton(self.trace.color)
        self._color_btn.color_changed.connect(self._on_color)
        lay.addWidget(self._color_btn)

        # Source label  dev_id / ch_id
        src = QLabel(f"{self.trace.device_id} / {self.trace.channel_id}")
        src.setObjectName("traceSource")
        src.setMinimumWidth(120)
        lay.addWidget(src)

        # Label override
        self._label_edit = QLineEdit(self.trace.label)
        self._label_edit.setPlaceholderText("Label…")
        self._label_edit.setObjectName("traceLabel")
        self._label_edit.textChanged.connect(self._on_label)
        lay.addWidget(self._label_edit, 1)

        # Line style
        self._style_cb = QComboBox()
        self._style_cb.addItems(["solid", "dash", "dot"])
        self._style_cb.setCurrentText(self.trace.style)
        self._style_cb.setObjectName("traceStyleCb")
        self._style_cb.setFixedWidth(64)
        self._style_cb.currentTextChanged.connect(self._on_style)
        lay.addWidget(self._style_cb)

        # Width
        self._width_spin = QDoubleSpinBox()
        self._width_spin.setRange(0.5, 6.0)
        self._width_spin.setSingleStep(0.5)
        self._width_spin.setValue(self.trace.width)
        self._width_spin.setFixedWidth(56)
        self._width_spin.setObjectName("traceWidthSpin")
        self._width_spin.valueChanged.connect(self._on_width)
        lay.addWidget(self._width_spin)

        # Remove button
        rm = QToolButton()
        rm.setText("✕")
        rm.setObjectName("traceRemoveBtn")
        rm.setFixedSize(22, 22)
        rm.clicked.connect(lambda: self.removed.emit(self))
        lay.addWidget(rm)

    def _on_vis(self, v):   self.trace.visible = v;         self.changed.emit()
    def _on_color(self, c): self.trace.color   = c;         self.changed.emit()
    def _on_label(self, t): self.trace.label   = t;         self.changed.emit()
    def _on_style(self, s): self.trace.style   = s;         self.changed.emit()
    def _on_width(self, w): self.trace.width   = w;         self.changed.emit()


# ══════════════════════════════════════════════════════════════════════════════
#  Plot block widget  (one collapsible section per PlotConfig)
# ══════════════════════════════════════════════════════════════════════════════

class PlotBlock(QFrame):
    removed  = pyqtSignal(object)   # self
    changed  = pyqtSignal()
    move_up  = pyqtSignal(object)
    move_dn  = pyqtSignal(object)

    def __init__(self, plot_cfg: PlotConfig, registry: DeviceRegistry, index: int):
        super().__init__()
        self.plot_cfg = plot_cfg
        self.registry = registry
        self.index    = index
        self.setObjectName("plotBlock")
        self._trace_rows: list[TraceRow] = []
        self._build()

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

        # ── Header bar ──────────────────────────────────────────────
        hdr = QWidget(); hdr.setObjectName("plotBlockHeader"); hdr.setFixedHeight(32)
        hdr_lay = QHBoxLayout(hdr); hdr_lay.setContentsMargins(8, 0, 6, 0)

        self._title_edit = QLineEdit(self.plot_cfg.title)
        self._title_edit.setObjectName("plotBlockTitle")
        self._title_edit.textChanged.connect(self._on_title)
        hdr_lay.addWidget(self._title_edit, 1)

        up_btn = QToolButton(); up_btn.setText("▲"); up_btn.setObjectName("plotMoveBtn")
        up_btn.setFixedSize(22, 22); up_btn.clicked.connect(lambda: self.move_up.emit(self))
        dn_btn = QToolButton(); dn_btn.setText("▼"); dn_btn.setObjectName("plotMoveBtn")
        dn_btn.setFixedSize(22, 22); dn_btn.clicked.connect(lambda: self.move_dn.emit(self))
        rm_btn = QToolButton(); rm_btn.setText("✕"); rm_btn.setObjectName("plotRemoveBtn")
        rm_btn.setFixedSize(22, 22); rm_btn.clicked.connect(lambda: self.removed.emit(self))

        for b in (up_btn, dn_btn, rm_btn): hdr_lay.addWidget(b)
        outer.addWidget(hdr)

        # ── Settings row ────────────────────────────────────────────
        settings = QWidget(); settings.setObjectName("plotBlockSettings")
        s_lay = QHBoxLayout(settings); s_lay.setContentsMargins(8, 6, 8, 6)

        s_lay.addWidget(QLabel("Y Label:"))
        self._ylabel_edit = QLineEdit(self.plot_cfg.y_label)
        self._ylabel_edit.setPlaceholderText("e.g. Voltage (V)")
        self._ylabel_edit.setObjectName("plotYLabelEdit")
        self._ylabel_edit.textChanged.connect(self._on_ylabel)
        s_lay.addWidget(self._ylabel_edit)

        self._auto_chk = QCheckBox("Auto Y")
        self._auto_chk.setChecked(self.plot_cfg.y_auto)
        self._auto_chk.toggled.connect(self._on_auto)
        s_lay.addWidget(self._auto_chk)

        s_lay.addWidget(QLabel("Min:"))
        self._ymin = QDoubleSpinBox()
        self._ymin.setRange(-1e9, 1e9); self._ymin.setValue(self.plot_cfg.y_min)
        self._ymin.setFixedWidth(72); self._ymin.setEnabled(not self.plot_cfg.y_auto)
        self._ymin.valueChanged.connect(self._on_ymin)
        s_lay.addWidget(self._ymin)

        s_lay.addWidget(QLabel("Max:"))
        self._ymax = QDoubleSpinBox()
        self._ymax.setRange(-1e9, 1e9); self._ymax.setValue(self.plot_cfg.y_max)
        self._ymax.setFixedWidth(72); self._ymax.setEnabled(not self.plot_cfg.y_auto)
        self._ymax.valueChanged.connect(self._on_ymax)
        s_lay.addWidget(self._ymax)

        self._grid_chk = QCheckBox("Grid")
        self._grid_chk.setChecked(self.plot_cfg.grid)
        self._grid_chk.toggled.connect(self._on_grid)
        s_lay.addWidget(self._grid_chk)

        s_lay.addStretch()
        outer.addWidget(settings)

        # ── Traces container ─────────────────────────────────────────
        self._traces_widget = QWidget()
        self._traces_widget.setObjectName("plotBlockTraces")
        self._traces_layout = QVBoxLayout(self._traces_widget)
        self._traces_layout.setContentsMargins(4, 2, 4, 4)
        self._traces_layout.setSpacing(2)

        for tr in self.plot_cfg.traces:
            self._add_trace_row(tr)

        # ── Add trace button ─────────────────────────────────────────
        add_row = QHBoxLayout()
        self._ch_picker = self._build_channel_picker()
        add_row.addWidget(self._ch_picker, 1)
        add_tr_btn = QPushButton("＋ Add Trace")
        add_tr_btn.setObjectName("addTraceBtn")
        add_tr_btn.clicked.connect(self._on_add_trace)
        add_row.addWidget(add_tr_btn)
        self._traces_layout.addLayout(add_row)

        outer.addWidget(self._traces_widget)

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

    def _build_channel_picker(self) -> QComboBox:
        cb = QComboBox()
        cb.setObjectName("channelPickerCb")
        cb.setPlaceholderText("Select channel…")
        for dev in self.registry.all_instances():
            for ch in dev.info.channels:
                label = f"{dev.info.device_id}  /  {ch.channel_id}  ({ch.name})"
                cb.addItem(label, userData=(dev.info.device_id, ch.channel_id,
                                            ch.name, ch.color))
        return cb

    def _add_trace_row(self, trace: ChannelTrace):
        row = TraceRow(trace, self.registry)
        row.removed.connect(self._remove_trace_row)
        row.changed.connect(self.changed)
        self._trace_rows.append(row)
        # Insert before the add-row (last item)
        idx = self._traces_layout.count() - 1
        self._traces_layout.insertWidget(idx, row)

    def _remove_trace_row(self, row: TraceRow):
        self.plot_cfg.traces.remove(row.trace)
        self._trace_rows.remove(row)
        self._traces_layout.removeWidget(row)
        row.deleteLater()
        self.changed.emit()

    def _on_add_trace(self):
        data = self._ch_picker.currentData()
        if data is None:
            return
        dev_id, ch_id, ch_name, ch_color = data
        tr = ChannelTrace(device_id=dev_id, channel_id=ch_id,
                          label=ch_name, color=ch_color)
        self.plot_cfg.traces.append(tr)
        self._add_trace_row(tr)
        self.changed.emit()

    # ── Config slots ─────────────────────────────────────────────────────

    def _on_title(self, t):  self.plot_cfg.title   = t;     self.changed.emit()
    def _on_ylabel(self, t): self.plot_cfg.y_label  = t;    self.changed.emit()
    def _on_auto(self, v):
        self.plot_cfg.y_auto = v
        self._ymin.setEnabled(not v); self._ymax.setEnabled(not v)
        self.changed.emit()
    def _on_ymin(self, v):   self.plot_cfg.y_min   = v;     self.changed.emit()
    def _on_ymax(self, v):   self.plot_cfg.y_max   = v;     self.changed.emit()
    def _on_grid(self, v):   self.plot_cfg.grid     = v;    self.changed.emit()


# ══════════════════════════════════════════════════════════════════════════════
#  PlotConfigWindow  — the main floating builder
# ══════════════════════════════════════════════════════════════════════════════

class PlotConfigWindow(QWidget):
    """
    Floating plot configuration builder.

    Emits `config_applied(ChartConfig)` when the user clicks Apply or
    when live-preview is enabled and any setting changes.
    """

    config_applied = pyqtSignal(object)   # ChartConfig
    closed         = pyqtSignal()

    def __init__(self, registry: DeviceRegistry,
                 current_config: Optional[ChartConfig] = None,
                 parent=None):
        super().__init__(
            parent,
            Qt.WindowType.Window | Qt.WindowType.Tool,
        )
        self.registry = registry
        self.cfg      = deepcopy(current_config) if current_config \
                        else default_config(registry)

        self.setWindowTitle("Plot Configuration")
        self.setMinimumSize(680, 560)
        self.resize(760, 640)
        self._blocks: list[PlotBlock] = []
        self._build()
        self._populate()

    # ── Build shell ───────────────────────────────────────────────────────

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

        # ── Top bar ──────────────────────────────────────────────────
        topbar = QWidget(); topbar.setObjectName("cfgTopBar"); topbar.setFixedHeight(44)
        tb_lay = QHBoxLayout(topbar); tb_lay.setContentsMargins(12, 0, 12, 0)

        title = QLabel("📈  PLOT CONFIGURATION")
        title.setObjectName("cfgTopBarTitle")
        tb_lay.addWidget(title, 1)

        self._live_chk = QCheckBox("Live preview")
        self._live_chk.setChecked(False)
        self._live_chk.setObjectName("cfgLiveChk")
        tb_lay.addWidget(self._live_chk)

        root.addWidget(topbar)

        # ── Global settings strip ────────────────────────────────────
        glob = QWidget(); glob.setObjectName("cfgGlobalBar")
        g_lay = QHBoxLayout(glob); g_lay.setContentsMargins(12, 6, 12, 6)

        g_lay.addWidget(QLabel("Time window:"))
        self._win_spin = QDoubleSpinBox()
        self._win_spin.setRange(1, 3600); self._win_spin.setSuffix(" s")
        self._win_spin.setValue(self.cfg.time_window_s)
        self._win_spin.setObjectName("cfgGlobalSpin")
        self._win_spin.valueChanged.connect(self._on_global_change)
        g_lay.addWidget(self._win_spin)
        g_lay.addSpacing(16)

        self._linkx_chk = QCheckBox("Link X axes")
        self._linkx_chk.setChecked(self.cfg.link_x)
        self._linkx_chk.setObjectName("cfgLiveChk")
        self._linkx_chk.toggled.connect(self._on_global_change)
        g_lay.addWidget(self._linkx_chk)

        self._legend_chk = QCheckBox("Show legend")
        self._legend_chk.setChecked(self.cfg.show_legend)
        self._legend_chk.setObjectName("cfgLiveChk")
        self._legend_chk.toggled.connect(self._on_global_change)
        g_lay.addWidget(self._legend_chk)

        g_lay.addStretch()

        add_plot_btn = QPushButton("＋  Add Plot Row")
        add_plot_btn.setObjectName("addDeviceButton")
        add_plot_btn.clicked.connect(self._add_plot)
        g_lay.addWidget(add_plot_btn)

        root.addWidget(glob)

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

        # ── Scrollable plot blocks ───────────────────────────────────
        scroll = QScrollArea()
        scroll.setWidgetResizable(True)
        scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarPolicy.ScrollBarAlwaysOff)
        scroll.setObjectName("deviceScroll")

        self._container  = QWidget()
        self._blocks_lay = QVBoxLayout(self._container)
        self._blocks_lay.setContentsMargins(10, 10, 10, 10)
        self._blocks_lay.setSpacing(12)
        self._blocks_lay.addStretch()

        scroll.setWidget(self._container)
        root.addWidget(scroll, 1)

        # ── Bottom action bar ────────────────────────────────────────
        btm = QWidget(); btm.setObjectName("cfgBottomBar")
        b_lay = QHBoxLayout(btm); b_lay.setContentsMargins(12, 8, 12, 8)

        self._reset_btn = QPushButton("⟳  Reset to Default")
        self._reset_btn.setObjectName("configButton")
        self._reset_btn.clicked.connect(self._reset_defaults)
        b_lay.addWidget(self._reset_btn)

        self._export_btn = QPushButton("↓  Export JSON")
        self._export_btn.setObjectName("configButton")
        self._export_btn.clicked.connect(self._export_json)
        b_lay.addWidget(self._export_btn)

        self._import_btn = QPushButton("↑  Import JSON")
        self._import_btn.setObjectName("configButton")
        self._import_btn.clicked.connect(self._import_json)
        b_lay.addWidget(self._import_btn)

        b_lay.addStretch()

        self._apply_btn = QPushButton("✓  Apply")
        self._apply_btn.setObjectName("applyButton")
        self._apply_btn.clicked.connect(self._apply)
        b_lay.addWidget(self._apply_btn)

        root.addWidget(btm)

    # ── Populate from config ───────────────────────────────────────────────

    def _populate(self):
        for blk in self._blocks:
            self._blocks_lay.removeWidget(blk)
            blk.deleteLater()
        self._blocks.clear()

        for i, p in enumerate(self.cfg.plots):
            self._insert_block(p, i)

    def _insert_block(self, plot_cfg: PlotConfig, index: int):
        blk = PlotBlock(plot_cfg, self.registry, index)
        blk.removed.connect(self._remove_block)
        blk.changed.connect(self._on_block_changed)
        blk.move_up.connect(self._move_block_up)
        blk.move_dn.connect(self._move_block_dn)
        self._blocks.append(blk)
        self._blocks_lay.insertWidget(self._blocks_lay.count() - 1, blk)

    # ── Block management ──────────────────────────────────────────────────

    def _add_plot(self):
        p = PlotConfig(title=f"Plot {len(self.cfg.plots) + 1}")
        self.cfg.plots.append(p)
        self._insert_block(p, len(self._blocks))
        self._maybe_live()

    def _remove_block(self, blk: PlotBlock):
        if blk.plot_cfg in self.cfg.plots:
            self.cfg.plots.remove(blk.plot_cfg)
        self._blocks.remove(blk)
        self._blocks_lay.removeWidget(blk)
        blk.deleteLater()
        self._maybe_live()

    def _move_block_up(self, blk: PlotBlock):
        i = self._blocks.index(blk)
        if i == 0: return
        self.cfg.plots.insert(i - 1, self.cfg.plots.pop(i))
        self._blocks.insert(i - 1, self._blocks.pop(i))
        self._blocks_lay.removeWidget(blk)
        self._blocks_lay.insertWidget(i - 1, blk)
        self._maybe_live()

    def _move_block_dn(self, blk: PlotBlock):
        i = self._blocks.index(blk)
        if i >= len(self._blocks) - 1: return
        self.cfg.plots.insert(i + 1, self.cfg.plots.pop(i))
        self._blocks.insert(i + 1, self._blocks.pop(i))
        self._blocks_lay.removeWidget(blk)
        self._blocks_lay.insertWidget(i + 1, blk)
        self._maybe_live()

    # ── Global settings ───────────────────────────────────────────────────

    def _on_global_change(self):
        self.cfg.time_window_s = self._win_spin.value()
        self.cfg.link_x        = self._linkx_chk.isChecked()
        self.cfg.show_legend   = self._legend_chk.isChecked()
        self._maybe_live()

    def _on_block_changed(self):
        self._maybe_live()

    def _maybe_live(self):
        if self._live_chk.isChecked():
            self.config_applied.emit(deepcopy(self.cfg))

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

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

    def _reset_defaults(self):
        self.cfg = default_config(self.registry)
        self._win_spin.setValue(self.cfg.time_window_s)
        self._populate()
        self._maybe_live()

    def _export_json(self):
        from PyQt6.QtWidgets import QFileDialog
        path, _ = QFileDialog.getSaveFileName(
            self, "Export Plot Config", "plot_config.json",
            "JSON Files (*.json)"
        )
        if path:
            with open(path, "w") as f:
                f.write(self.cfg.to_json())

    def _import_json(self):
        from PyQt6.QtWidgets import QFileDialog
        path, _ = QFileDialog.getOpenFileName(
            self, "Import Plot Config", "", "JSON Files (*.json)"
        )
        if path:
            try:
                with open(path) as f:
                    self.cfg = ChartConfig.from_json(f.read())
                self._win_spin.setValue(self.cfg.time_window_s)
                self._populate()
                self._maybe_live()
            except Exception as e:
                from PyQt6.QtWidgets import QMessageBox
                QMessageBox.critical(self, "Import failed", str(e))

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

    def update_registry(self):
        """Call after devices are added/removed to refresh channel pickers."""
        self._populate()