path: root/ui/strip_chart.py

"""
ui/strip_chart.py  —  Config-driven live chart.

Consumes LayoutConfig from plot_window.
Supports stacked, side-by-side, and grid layouts.
X axis can be time or any channel (e.g. stress vs strain).
"""

import numpy as np
from typing import Dict, List, Optional, Tuple
from collections import deque

from PyQt6.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QLabel,
    QDoubleSpinBox, QPushButton, QSizePolicy,
)
from PyQt6.QtCore import Qt, pyqtSlot, QTimer
from PyQt6.QtGui import QFont

try:
    import pyqtgraph as pg
    _HAS_PG = True
except ImportError:
    _HAS_PG = False

_THEME_DARK  = {"bg": "#0b0e13", "fg": "#475569", "grid_alpha": 0.12, "legend_brush": "#161d2e", "legend_pen": "#2a3558", "tick_color": "#64748b"}
_THEME_LIGHT = {"bg": "#ffffff", "fg": "#334155", "grid_alpha": 0.18, "legend_brush": "#f8fafc", "legend_pen": "#cbd5e1", "tick_color": "#64748b"}
_current_theme = _THEME_DARK

def set_chart_theme(theme: str):
    global _current_theme
    _current_theme = _THEME_DARK if theme == "dark" else _THEME_LIGHT
    if _HAS_PG:
        pg.setConfigOptions(antialias=True,
                            background=_current_theme["bg"],
                            foreground=_current_theme["fg"])

# Apply dark theme as default
if _HAS_PG:
    pg.setConfigOptions(antialias=True, background=_THEME_DARK["bg"], foreground=_THEME_DARK["fg"])

from devices.device_registry import DeviceRegistry
from core.acquisition import AcquisitionEngine
from core.signal_processor import SignalProcessor
from ui.windows.plot_window import (LayoutConfig, PaneSpec, TraceSpec,
                                    build_default_layout, tree_to_grid, _tree_grid_size)

_STYLES = {
    "solid": Qt.PenStyle.SolidLine,
    "dash":  Qt.PenStyle.DashLine,
    "dot":   Qt.PenStyle.DotLine,
}


def _windowed(times_dq, values_dq, seconds: float):
    """Last `seconds` of a (times, values) deque pair, walking in from the
    newest sample only — same trick as ChannelBuffer.window(), needed here
    for derived-channel buffers which aren't wrapped in ChannelBuffer."""
    if not times_dq:
        return [], []
    cutoff = times_dq[-1] - seconds
    out_t, out_v = [], []
    for t, v in zip(reversed(times_dq), reversed(values_dq)):
        if t < cutoff:
            break
        out_t.append(t); out_v.append(v)
    out_t.reverse(); out_v.reverse()
    return out_t, out_v


class StripChartWidget(QWidget):
    def __init__(self, engine: AcquisitionEngine, registry: DeviceRegistry,
                 processor: SignalProcessor):
        super().__init__()
        self.engine    = engine
        self.registry  = registry
        self.processor = processor
        self._paused   = False
        self._cfg: Optional[LayoutConfig] = None
        # (dev_id, ch_id) → [{curve, plot, spec}]
        self._curves: Dict[Tuple, List[dict]] = {}
        # Time-axis plots, scrolled continuously by _scroll_timer rather than
        # only when a new sample lands — decouples the visible scroll rate
        # from the (much slower) poll rate so it doesn't look stepped.
        self._time_plots: List = []
        self._build()
        self.apply_layout(build_default_layout(registry, processor))

        self._scroll_timer = QTimer(self)
        self._scroll_timer.setInterval(33)   # ~30 fps
        self._scroll_timer.timeout.connect(self._on_scroll_tick)
        self._scroll_timer.start()

    def _build(self):
        lay = QVBoxLayout(self); lay.setContentsMargins(6,6,6,4); lay.setSpacing(4)
        ctrl = QHBoxLayout()
        ctrl.addWidget(QLabel("Window:"))
        self._win = QDoubleSpinBox(); self._win.setRange(1,3600)
        self._win.setSuffix(" s"); self._win.setValue(30.0)
        self._win.valueChanged.connect(lambda v: self._cfg and setattr(self._cfg,"time_window_s",v))
        ctrl.addWidget(self._win); ctrl.addStretch()
        self._pause = QPushButton("⏸  Pause"); self._pause.setCheckable(True)
        self._pause.setObjectName("pauseButton"); self._pause.toggled.connect(self._on_pause)
        ctrl.addWidget(self._pause); lay.addLayout(ctrl)
        if _HAS_PG:
            self._gw = pg.GraphicsLayoutWidget()
            self._gw.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
            lay.addWidget(self._gw)
        else:
            lay.addWidget(QLabel("⚠  pip install pyqtgraph",
                                  alignment=Qt.AlignmentFlag.AlignCenter))

    def apply_layout(self, cfg: LayoutConfig):
        self._cfg = cfg
        self._win.setValue(cfg.time_window_s)
        if not _HAS_PG: return
        self._gw.clear(); self._curves.clear(); self._time_plots.clear()
        if not cfg.panes: return

        mf = QFont("IBM Plex Mono", 8)
        n  = len(cfg.panes)

        # Determine row/col/rowspan/colspan for each pane
        if cfg.layout_mode == "free" and cfg.tree is not None:
            gs = _tree_grid_size(cfg.tree)
            grid_pos: dict = {}
            tree_to_grid(cfg.tree, 0, 0, gs, gs, grid_pos)
            # grid_pos[pane_idx] = (row, col, rowspan, colspan)
        elif cfg.layout_mode == "sidebyside":
            grid_pos = {i: (0, i, 1, 1) for i in range(n)}
        elif cfg.layout_mode == "grid":
            cols = max(1, cfg.grid_cols)
            grid_pos = {i: (i // cols, i % cols, 1, 1) for i in range(n)}
        else:  # stacked (legacy default)
            grid_pos = {i: (i, 0, 1, 1) for i in range(n)}

        ref_plot = None
        for idx, spec in enumerate(cfg.panes):
            if idx not in grid_pos:
                continue
            row, col, rowspan, colspan = grid_pos[idx]
            plot = self._gw.addPlot(row=row, col=col, rowspan=rowspan, colspan=colspan)
            plot.setLabel("left",   spec.y_label or spec.title)
            plot.setLabel("bottom", spec.x_label or "Elapsed (s)")
            plot.showGrid(x=spec.grid, y=spec.grid, alpha=0.12)
            plot.getAxis("left").setStyle(tickFont=mf)
            plot.getAxis("bottom").setStyle(tickFont=mf)

            # Hide bottom axis labels for non-bottom panes in stacked legacy mode
            if cfg.layout_mode == "stacked" and idx < n - 1:
                plot.getAxis("bottom").setStyle(showValues=False)
                plot.getAxis("bottom").setHeight(0)

            if cfg.link_x and ref_plot and spec.x_source == "time":
                plot.setXLink(ref_plot)
            if spec.x_source == "time":
                ref_plot = plot
                self._time_plots.append(plot)

            if cfg.show_legend:
                t = _current_theme
                plot.addLegend(offset=(5,5), labelTextColor=t["tick_color"],
                               brush=pg.mkBrush(t["legend_brush"]),
                               pen=pg.mkPen(t["legend_pen"]))

            if not spec.y_auto:
                plot.setYRange(spec.y_min, spec.y_max)
            else:
                plot.enableAutoRange(axis="y")

            self._gw.ci.layout.setColumnStretchFactor(col, spec.weight)
            self._gw.ci.layout.setRowStretchFactor(row, spec.weight)

            for tr in spec.traces:
                if not tr.visible: continue
                pen   = pg.mkPen(color=tr.color, width=tr.width,
                                  style=_STYLES.get(tr.style, Qt.PenStyle.SolidLine))
                curve = plot.plot([], [], pen=pen, name=tr.label or tr.channel_id)
                key   = (tr.device_id, tr.channel_id)
                self._curves.setdefault(key, []).append(
                    {"curve": curve, "plot": plot, "spec": spec, "pane": spec})

    def set_theme(self, theme: str):
        """Hot-swap pyqtgraph colours when the user changes theme."""
        set_chart_theme(theme)
        if _HAS_PG and self._gw is not None:
            self._gw.setBackground(_current_theme["bg"])
        # Rebuild plots so legend/grid colours update
        self.apply_layout(self._cfg or build_default_layout(self.registry, self.processor))

    def refresh(self):
        self.apply_layout(self._cfg or build_default_layout(self.registry, self.processor))

    def on_channel_enabled_changed(self, dev_id: str, ch_id: str, enabled: bool):
        key = (dev_id, ch_id)
        for e in self._curves.get(key, []):
            e["curve"].setVisible(enabled)

    @pyqtSlot(str, str, float, float)
    def on_new_data(self, device_id: str, channel_id: str, ts: float, val: float):
        if self._paused or not _HAS_PG or not self._cfg: return
        key     = (device_id, channel_id)
        entries = self._curves.get(key)
        if not entries: return
        window  = self._cfg.time_window_s

        # Get Y data
        if device_id == "derived":
            bufs = self.processor.get_derived_buffer(channel_id)
            if bufs is None or len(bufs[0]) < 2: return
            ts_l, vs_l = _windowed(bufs[0], bufs[1], window)
            if not ts_l: return
        else:
            buf = self.engine.get_buffer(device_id, channel_id)
            if buf is None or len(buf) < 2: return
            ts_l, vs_l = buf.window(window)
            if not ts_l: return

        for e in entries:
            spec = e["pane"]
            # Determine X axis data
            if spec.x_source == "time":
                x_arr = np.array(ts_l, dtype=np.float64)
                y_arr = np.array(vs_l, dtype=np.float64)
                e["curve"].setData(x_arr, y_arr)
                # X-range scrolling is handled by _scroll_timer, continuously —
                # not here, so the window doesn't jump only on data arrival.
            else:
                # X = another channel's processed values
                parts = spec.x_source.split("/")
                if len(parts) == 2:
                    x_dev, x_ch = parts
                    if x_dev == "derived":
                        xbufs = self.processor.get_derived_buffer(x_ch)
                        if xbufs and len(xbufs[0]) >= 2:
                            x_arr = np.array(list(xbufs[1]), dtype=np.float64)
                            y_arr = np.array(vs_l, dtype=np.float64)
                            mn = min(len(x_arr), len(y_arr))
                            e["curve"].setData(x_arr[-mn:], y_arr[-mn:])
                    else:
                        xbuf = self.engine.get_buffer(x_dev, x_ch)
                        if xbuf and len(xbuf) >= 2:
                            _, xvs = xbuf.all()
                            x_arr = np.array(xvs, dtype=np.float64)
                            y_arr = np.array(vs_l, dtype=np.float64)
                            mn = min(len(x_arr), len(y_arr))
                            e["curve"].setData(x_arr[-mn:], y_arr[-mn:])

            if spec.y_auto:
                e["plot"].enableAutoRange(axis="y")

    def _on_pause(self, c: bool):
        self._paused = c
        self._pause.setText("▶  Resume" if c else "⏸  Pause")

    def _on_scroll_tick(self):
        """Scroll time-axis plots at a steady UI rate, independent of the
        (much slower) poll rate — keeps the window sliding smoothly instead
        of jumping once per data tick."""
        if self._paused or not _HAS_PG or not self._cfg or not self.engine.is_running():
            return
        now    = self.engine.elapsed_now()
        window = self._cfg.time_window_s
        for plot in self._time_plots:
            plot.setXRange(now - window, now, padding=0)