we_monitor/gui_client.py

#!/usr/bin/env python3
"""
WeConnect GUI client.

Connects to the collect.py TCP push server and displays live data.

Tabs
----
Connector  — server address, connect / disconnect, status
Dashboard  — live tree view of the latest snapshot
Plot       — time-series charts of all physical {value, unit} fields
"""

import json
import socket
import sys
import threading
from collections import deque
from datetime import datetime

from PyQt5.QtCore import QObject, Qt, pyqtSignal
from PyQt5.QtGui import QColor, QFont
from PyQt5.QtWidgets import (
    QApplication, QFormLayout, QGroupBox, QHBoxLayout, QLabel,
    QLineEdit, QMainWindow, QPushButton, QScrollArea, QSizePolicy,
    QSpinBox, QStatusBar, QTabWidget, QTreeWidget, QTreeWidgetItem,
    QVBoxLayout, QWidget,
)

import matplotlib
matplotlib.use("Qt5Agg")
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure

MAX_PLOT_POINTS = 200
PLOT_COLS       = 2
PLOT_ROW_INCHES = 3.0


# ── TCP reader ────────────────────────────────────────────────────────────────

class TcpReader(QObject):
    """Background thread that emits decoded JSON messages via Qt signals."""

    message        = pyqtSignal(dict)
    status_changed = pyqtSignal(str)   # emitted only on unexpected loss

    def __init__(self):
        super().__init__()
        self._sock    = None
        self._running = False

    def connect_to(self, host: str, port: int) -> str | None:
        """Open connection; returns error string or None on success."""
        try:
            self._sock = socket.create_connection((host, port), timeout=5)
        except OSError as exc:
            return str(exc)
        self._running = True
        threading.Thread(target=self._read_loop, daemon=True).start()
        return None

    def disconnect(self):
        self._running = False
        if self._sock:
            try:
                self._sock.close()
            except OSError:
                pass
            self._sock = None

    def _read_loop(self):
        buf = ""
        try:
            while self._running:
                chunk = self._sock.recv(4096)
                if not chunk:
                    break
                buf += chunk.decode(errors="replace")
                while "\n" in buf:
                    line, buf = buf.split("\n", 1)
                    line = line.strip()
                    if line:
                        try:
                            self.message.emit(json.loads(line))
                        except json.JSONDecodeError:
                            pass
        except OSError:
            pass
        if self._running:
            self._running = False
            self.status_changed.emit("Connection lost")


# ── helpers ───────────────────────────────────────────────────────────────────

def _is_phys(v) -> bool:
    return isinstance(v, dict) and "value" in v and "unit" in v


def _extract_phys(data: dict, prefix: str = "") -> dict[str, dict]:
    """Return a flat dict of path → {value, unit} for every physical leaf."""
    result = {}
    for k, v in data.items():
        if k == "ts":
            continue
        path = f"{prefix}.{k}" if prefix else k
        if _is_phys(v):
            result[path] = v
        elif isinstance(v, dict):
            result.update(_extract_phys(v, path))
    return result


# ── Connector tab ─────────────────────────────────────────────────────────────

class ConnectorTab(QWidget):
    connect_requested    = pyqtSignal(str, int)
    disconnect_requested = pyqtSignal()

    def __init__(self):
        super().__init__()
        layout = QVBoxLayout(self)
        layout.setAlignment(Qt.AlignTop)

        grp  = QGroupBox("Server")
        form = QFormLayout(grp)

        self._host = QLineEdit("127.0.0.1")
        self._port = QSpinBox()
        self._port.setRange(1, 65535)
        self._port.setValue(9999)
        form.addRow("Host:", self._host)
        form.addRow("Port:", self._port)

        btn_row = QHBoxLayout()
        self._btn = QPushButton("Connect")
        self._btn.setFixedWidth(110)
        self._btn.clicked.connect(self._toggle)
        btn_row.addWidget(self._btn)
        btn_row.addStretch()
        form.addRow(btn_row)

        self._status = QLabel("Disconnected")
        self._status.setStyleSheet("color: gray;")
        form.addRow("Status:", self._status)

        layout.addWidget(grp)

    def _toggle(self):
        if self._btn.text() == "Connect":
            self.connect_requested.emit(self._host.text(), self._port.value())
        else:
            self.disconnect_requested.emit()

    def set_connected(self, host: str, port: int):
        self._btn.setText("Disconnect")
        self._host.setEnabled(False)
        self._port.setEnabled(False)
        self._status.setText(f"Connected  →  {host}:{port}")
        self._status.setStyleSheet("color: green; font-weight: bold;")

    def set_disconnected(self, reason: str = ""):
        self._btn.setText("Connect")
        self._host.setEnabled(True)
        self._port.setEnabled(True)
        if reason:
            self._status.setText(f"Disconnected  ({reason})")
            self._status.setStyleSheet("color: red;")
        else:
            self._status.setText("Disconnected")
            self._status.setStyleSheet("color: gray;")


# ── Dashboard tab ─────────────────────────────────────────────────────────────

class DashboardTab(QWidget):
    def __init__(self):
        super().__init__()
        layout = QVBoxLayout(self)

        self._ts = QLabel("Last update: —")
        self._ts.setAlignment(Qt.AlignRight)
        self._ts.setStyleSheet("color: gray; font-size: 11px;")
        layout.addWidget(self._ts)

        self._tree = QTreeWidget()
        self._tree.setColumnCount(3)
        self._tree.setHeaderLabels(["Field", "Value", "Unit"])
        self._tree.setColumnWidth(0, 320)
        self._tree.setColumnWidth(1, 180)
        self._tree.setColumnWidth(2, 80)
        self._tree.setAlternatingRowColors(True)
        self._tree.setRootIsDecorated(True)
        layout.addWidget(self._tree)

    def update(self, snapshot: dict):
        self._ts.setText(f"Last update:  {snapshot.get('ts', '—')}")

        expanded = self._expanded_paths()
        self._tree.clear()

        bold = QFont()
        bold.setBold(True)
        domain_color = QColor("#1a4d99")

        for domain, domain_data in snapshot.items():
            if domain == "ts" or not isinstance(domain_data, dict):
                continue
            d_item = QTreeWidgetItem([domain.upper()])
            d_item.setFont(0, bold)
            d_item.setForeground(0, domain_color)
            for obj_name, fields in domain_data.items():
                if not isinstance(fields, dict):
                    continue
                obj_item = QTreeWidgetItem([obj_name])
                obj_item.setFont(0, bold)
                self._add_fields(obj_item, fields, f"{domain}/{obj_name}")
                d_item.addChild(obj_item)
            self._tree.addTopLevelItem(d_item)

        self._tree.expandAll()
        self._restore_expanded(expanded)

    def _add_fields(self, parent: QTreeWidgetItem, data: dict, path: str):
        for k, v in data.items():
            child_path = f"{path}/{k}"
            if _is_phys(v):
                item = QTreeWidgetItem([k, str(v["value"]), v["unit"]])
                item.setTextAlignment(1, Qt.AlignRight | Qt.AlignVCenter)
                parent.addChild(item)
            elif isinstance(v, dict):
                node = QTreeWidgetItem([k])
                self._add_fields(node, v, child_path)
                parent.addChild(node)
            else:
                item = QTreeWidgetItem([k, str(v), ""])
                item.setTextAlignment(1, Qt.AlignRight | Qt.AlignVCenter)
                parent.addChild(item)

    # ── preserve expand state across refreshes ──────────────────────────────

    def _expanded_paths(self) -> set[str]:
        paths = set()
        root = self._tree.invisibleRootItem()
        self._collect_expanded(root, "", paths)
        return paths

    def _collect_expanded(self, item, prefix, paths):
        for i in range(item.childCount()):
            child = item.child(i)
            path = f"{prefix}/{child.text(0)}"
            if child.isExpanded():
                paths.add(path)
            self._collect_expanded(child, path, paths)

    def _restore_expanded(self, paths: set[str]):
        root = self._tree.invisibleRootItem()
        self._apply_expanded(root, "", paths)

    def _apply_expanded(self, item, prefix, paths):
        for i in range(item.childCount()):
            child = item.child(i)
            path = f"{prefix}/{child.text(0)}"
            if path in paths:
                child.setExpanded(True)
            self._apply_expanded(child, path, paths)


# ── Plot tab ──────────────────────────────────────────────────────────────────

class PlotTab(QWidget):
    def __init__(self):
        super().__init__()
        outer = QVBoxLayout(self)
        outer.setContentsMargins(0, 0, 0, 0)

        scroll = QScrollArea()
        scroll.setWidgetResizable(True)
        outer.addWidget(scroll)

        container = QWidget()
        inner = QVBoxLayout(container)
        inner.setContentsMargins(4, 4, 4, 4)
        scroll.setWidget(container)

        self._fig    = Figure(tight_layout=True)
        self._canvas = FigureCanvas(self._fig)
        self._canvas.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        inner.addWidget(self._canvas)

        self._series: dict[str, deque]  = {}   # path → deque[(datetime, float)]
        self._units:  dict[str, str]    = {}   # path → unit string
        self._axes:   dict[str, object] = {}   # path → matplotlib Axes

    def update(self, snapshot: dict):
        ts_raw = snapshot.get("ts", "")
        try:
            ts = datetime.fromisoformat(ts_raw)
        except (ValueError, TypeError):
            ts = datetime.utcnow()

        phys = _extract_phys(snapshot)
        if not phys:
            return

        new_keys = [k for k in phys if k not in self._series]
        for k in new_keys:
            self._series[k] = deque(maxlen=MAX_PLOT_POINTS)

        for k, pv in phys.items():
            self._units[k] = pv.get("unit", "")
            try:
                self._series[k].append((ts, float(pv["value"])))
            except (TypeError, ValueError):
                pass

        if new_keys:
            self._rebuild_layout()

        self._redraw()

    def _rebuild_layout(self):
        keys = list(self._series.keys())
        n    = len(keys)
        cols = min(PLOT_COLS, n)
        rows = (n + cols - 1) // cols

        self._fig.clear()
        self._axes = {}

        height = max(3.0, rows * PLOT_ROW_INCHES)
        self._fig.set_size_inches(10, height)
        self._canvas.setMinimumHeight(int(height * self._fig.dpi))

        for i, k in enumerate(keys):
            ax = self._fig.add_subplot(rows, cols, i + 1)
            label = k.split(".")[-1]
            ax.set_title(label, fontsize=9, pad=4)
            ax.set_ylabel(self._units.get(k, ""), fontsize=8)
            ax.tick_params(axis="both", labelsize=7)
            ax.grid(True, linewidth=0.4, alpha=0.5)
            self._axes[k] = ax

    def _redraw(self):
        for k, ax in self._axes.items():
            series = self._series.get(k)
            if not series:
                continue
            times, values = zip(*series)
            ax.clear()
            ax.set_title(k.split(".")[-1], fontsize=9, pad=4)
            ax.set_ylabel(self._units.get(k, ""), fontsize=8)
            ax.tick_params(axis="both", labelsize=7)
            ax.grid(True, linewidth=0.4, alpha=0.5)
            ax.plot(times, values, linewidth=1.4, color="#1a6eb5")
            if len(times) > 1:
                ax.set_xlim(times[0], times[-1])
        self._fig.autofmt_xdate(rotation=25, ha="right")
        self._canvas.draw_idle()


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

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("WeConnect Monitor")
        self.resize(960, 700)

        self._reader = TcpReader()
        self._reader.message.connect(self._on_message)
        self._reader.status_changed.connect(self._on_connection_lost)

        self._connector = ConnectorTab()
        self._dashboard = DashboardTab()
        self._plot      = PlotTab()

        tabs = QTabWidget()
        tabs.addTab(self._connector, "Connector")
        tabs.addTab(self._dashboard, "Dashboard")
        tabs.addTab(self._plot,      "Plot")
        self.setCentralWidget(tabs)

        self._statusbar = QStatusBar()
        self.setStatusBar(self._statusbar)
        self._statusbar.showMessage("Disconnected")

        self._connector.connect_requested.connect(self._connect)
        self._connector.disconnect_requested.connect(self._disconnect)

    def _connect(self, host: str, port: int):
        err = self._reader.connect_to(host, port)
        if err:
            self._connector.set_disconnected(err)
            self._statusbar.showMessage(f"Connection failed: {err}")
        else:
            self._connector.set_connected(host, port)
            self._statusbar.showMessage(f"Connected to {host}:{port}")

    def _disconnect(self):
        self._reader.disconnect()
        self._connector.set_disconnected()
        self._statusbar.showMessage("Disconnected")

    def _on_message(self, data: dict):
        self._dashboard.update(data)
        self._plot.update(data)
        self._statusbar.showMessage(
            f"Last snapshot: {data.get('ts', '?')}"
        )

    def _on_connection_lost(self, reason: str):
        self._connector.set_disconnected(reason)
        self._statusbar.showMessage(reason)


# ── entry point ───────────────────────────────────────────────────────────────

def main():
    app = QApplication(sys.argv)
    app.setStyle("Fusion")
    win = MainWindow()
    win.show()
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()