mppt-testbench/testbench/gui_workers.py

"""Background worker thread for instrument I/O.

Keeps all VISA/serial communication off the GUI thread so tkinter
never freezes during instrument queries.
"""

from __future__ import annotations

import csv
import queue
import threading
import time
from collections import deque
from enum import Enum, auto

import serial

from testbench.stm32_link import (
    BroadcastData, BroadcastParser, PARAM_BY_ID, decode_param_bits,
)


class Cmd(Enum):
    """Commands sent from GUI to worker thread."""

    # Supply
    SET_VOLTAGE = auto()
    SET_CURRENT = auto()
    APPLY = auto()
    OUTPUT_ON = auto()
    OUTPUT_OFF = auto()
    SET_OVP = auto()
    SET_RISE_TIME = auto()
    SET_FALL_TIME = auto()

    # Load
    SET_MODE = auto()
    SET_MODE_VALUE = auto()
    LOAD_ON = auto()
    LOAD_OFF = auto()
    SET_SLEW_RISE = auto()
    SET_SLEW_FALL = auto()

    # Meter
    SETUP_ALL = auto()
    SET_WIRING = auto()
    SET_COUPLING = auto()
    SET_RESPONSE_SPEED = auto()
    SET_AVERAGING = auto()
    SET_VOLTAGE_RANGE = auto()
    SET_CURRENT_RANGE = auto()
    SET_VOLTAGE_AUTO = auto()
    SET_CURRENT_AUTO = auto()
    DEGAUSS = auto()

    # System
    SET_INTERVAL = auto()
    SAFE_OFF = auto()


class InstrumentWorker(threading.Thread):
    """Daemon thread that polls instruments and processes commands.

    Usage:
        worker = InstrumentWorker(bench)
        worker.start()
        worker.send(Cmd.SET_VOLTAGE, 48.0)
        data = worker.get_data()  # dict or None
        worker.stop()
    """

    def __init__(self, bench, interval: float = 1.0) -> None:
        super().__init__(daemon=True)
        self.bench = bench
        self.interval = interval
        self.cmd_queue: queue.Queue = queue.Queue()
        self.data_queue: queue.Queue = queue.Queue(maxsize=100)
        self._stop_event = threading.Event()

    def send(self, cmd: Cmd, *args) -> None:
        """Queue a command for execution on the worker thread."""
        self.cmd_queue.put((cmd, args))

    def get_data(self) -> dict | None:
        """Get the latest measurement data (non-blocking). Returns None if empty."""
        result = None
        # Drain queue, keep only the latest
        while True:
            try:
                result = self.data_queue.get_nowait()
            except queue.Empty:
                break
        return result

    def stop(self) -> None:
        """Signal the worker to stop."""
        self._stop_event.set()

    def run(self) -> None:
        """Main worker loop: process commands, then read instruments."""
        while not self._stop_event.is_set():
            # Process all pending commands
            while True:
                try:
                    cmd, args = self.cmd_queue.get_nowait()
                    self._execute(cmd, args)
                except queue.Empty:
                    break

            # Read all instruments
            try:
                data = self.bench.measure_all()
                data["_error"] = None
                data["_timestamp"] = time.time()
                # Query output states for GUI indicators
                try:
                    data["supply_on"] = self.bench.supply.get_output_state()
                except Exception:
                    data["supply_on"] = None
                try:
                    data["load_on"] = self.bench.load.get_load_state()
                except Exception:
                    data["load_on"] = None
                # Query meter channel ranges
                for ch in (5, 6):
                    try:
                        data[f"v_range_{ch}"] = self.bench.meter.get_voltage_range(ch).strip()
                    except Exception:
                        data[f"v_range_{ch}"] = None
                    try:
                        data[f"i_range_{ch}"] = self.bench.meter.get_current_range(ch).strip()
                    except Exception:
                        data[f"i_range_{ch}"] = None
            except Exception as e:
                data = {"_error": str(e), "_timestamp": time.time()}

            # Push to GUI (drop oldest if full)
            try:
                self.data_queue.put_nowait(data)
            except queue.Full:
                try:
                    self.data_queue.get_nowait()
                except queue.Empty:
                    pass
                self.data_queue.put_nowait(data)

            self._stop_event.wait(timeout=self.interval)

    def _execute(self, cmd: Cmd, args: tuple) -> None:
        """Execute a single command. Exceptions are swallowed and reported."""
        try:
            bench = self.bench
            match cmd:
                # Supply
                case Cmd.SET_VOLTAGE:
                    bench.supply.set_voltage(args[0])
                case Cmd.SET_CURRENT:
                    bench.supply.set_current(args[0])
                case Cmd.APPLY:
                    bench.supply.set_current(args[1])
                    bench.supply.set_voltage(args[0])
                case Cmd.OUTPUT_ON:
                    bench.supply.output_on()
                case Cmd.OUTPUT_OFF:
                    bench.supply.output_off()
                case Cmd.SET_OVP:
                    bench.supply.set_ovp_level(args[0])
                    bench.supply.set_ovp_state(True)
                case Cmd.SET_RISE_TIME:
                    bench.supply.set_rise_time(args[0])
                case Cmd.SET_FALL_TIME:
                    bench.supply.set_fall_time(args[0])

                # Load
                case Cmd.SET_MODE:
                    bench.load.set_mode(args[0])
                case Cmd.SET_MODE_VALUE:
                    mode, value = args[0], args[1]
                    if mode == "CC":
                        bench.load.set_cc_current(value)
                    elif mode == "CR":
                        bench.load.set_cr_resistance(value)
                    elif mode == "CV":
                        bench.load.set_cv_voltage(value)
                    elif mode == "CP":
                        bench.load.set_cp_power(value)
                case Cmd.LOAD_ON:
                    bench.load.load_on()
                case Cmd.LOAD_OFF:
                    bench.load.load_off()
                case Cmd.SET_SLEW_RISE:
                    bench.load.set_rise_slew(args[0])
                case Cmd.SET_SLEW_FALL:
                    bench.load.set_fall_slew(args[0])

                # Meter
                case Cmd.SETUP_ALL:
                    bench.setup_all()
                case Cmd.SET_WIRING:
                    bench.meter.set_wiring_mode(args[0])
                case Cmd.SET_COUPLING:
                    bench.meter.set_coupling(args[0], args[1])
                case Cmd.SET_RESPONSE_SPEED:
                    bench.meter.set_response_speed(args[0])
                case Cmd.SET_AVERAGING:
                    bench.meter.set_averaging(args[0], args[1] if len(args) > 1 else None)
                case Cmd.SET_VOLTAGE_RANGE:
                    bench.meter.set_voltage_auto(args[0], False)
                    bench.meter.set_voltage_range(args[0], args[1])
                case Cmd.SET_CURRENT_RANGE:
                    bench.meter.set_current_auto(args[0], False)
                    bench.meter.set_current_range(args[0], args[1])
                case Cmd.SET_VOLTAGE_AUTO:
                    bench.meter.set_voltage_auto(args[0], args[1])
                case Cmd.SET_CURRENT_AUTO:
                    bench.meter.set_current_auto(args[0], args[1])
                case Cmd.DEGAUSS:
                    channels = args[0] if args else [5, 6]
                    items = ",".join(f"I{ch}" for ch in channels)
                    bench.meter.write(f":DEMAg {items}")

                # System
                case Cmd.SET_INTERVAL:
                    self.interval = args[0]
                case Cmd.SAFE_OFF:
                    bench.safe_off()

        except Exception as e:
            # Push error to data queue so GUI can display it
            try:
                self.data_queue.put_nowait({
                    "_error": f"Command {cmd.name} failed: {e}",
                    "_timestamp": time.time(),
                })
            except queue.Full:
                pass


# ── STM32 broadcast reader ───────────────────────────────────────────

# All BroadcastData wire fields, in dataclass order (full-rate CSV columns).
TELEM_CSV_FIELDS = (
    "counter", "vin", "vout", "iin", "iout", "vfly", "etemp", "btemp",
    "vfly_integral", "vfly_avg_debug", "cc_output_f", "mppt_iref",
    "mppt_last_vin", "mppt_last_iin", "p_in", "p_out", "iout_slow",
    "last_tmp", "VREF", "vfly_correction", "cmp_outer", "cmp_inner",
    "vfly_ofs_applied", "ctrl_mode", "vfly_active", "status_flags",
    "fmac_sr", "fault_pc", "cfsr", "param_id", "param_type",
    "param_value", "pong", "sys_current_ma", "iin_avg_ma",
)
_TELEM_HEX_FIELDS = frozenset({"status_flags", "fmac_sr", "fault_pc", "cfsr"})


class STM32Worker(threading.Thread):
    """Daemon thread reading the STM32 broadcast stream.

    The board streams 114-byte frames continuously (100 Hz publishes, each
    repeated ~3-4x at line rate); this thread validates them, keeps the
    latest sample plus 20 s graph history, and owns the full-rate telemetry
    CSV (single writer). Auto-reconnects on serial errors.
    """

    HISTORY = 2200          # 100 Hz x 20 s + margin
    CSV_FLUSH_S = 1.0       # never flush per-row at 100 rows/s

    def __init__(self, port: str, baudrate: int = 460800) -> None:
        super().__init__(daemon=True)
        self.port = port
        self.baudrate = baudrate
        self.connected = False
        self._stop_event = threading.Event()
        self._tx_queue: queue.Queue = queue.Queue()

        self._lock = threading.Lock()
        self.latest: BroadcastData | None = None
        self.latest_wall: float = 0.0
        self._last_counter = -1
        self._t: deque = deque(maxlen=self.HISTORY)      # time.monotonic per fresh publish
        self._vfly: deque = deque(maxlen=self.HISTORY)   # mV
        self._corr: deque = deque(maxlen=self.HISTORY)   # vfly_correction, ticks
        self._ofs: deque = deque(maxlen=self.HISTORY)    # vfly_ofs_applied, ticks
        self._fresh_count = 0    # validated fresh publishes since last get_rates()
        self._sent_count = 0     # publishes the board sent (counter deltas)
        self.params: dict[int, float] = {}

        self._csv_lock = threading.Lock()
        self._csv_file = None
        self._csv_writer = None
        self._csv_last_flush = 0.0

    # ── Thread-safe API for the GUI ──────────────────────────────────

    def get_latest(self) -> tuple[BroadcastData | None, float]:
        """Latest validated sample and its wall-clock arrival time."""
        with self._lock:
            return self.latest, self.latest_wall

    def get_graph_snapshot(self, decimate: int = 4):
        """(t_mono, vfly_mV, corr, ofs) lists, decimated for display.

        Decimation is anchored at the end so the newest sample always shows.
        Full rate is still recorded to the CSV.
        """
        with self._lock:
            t, v = list(self._t), list(self._vfly)
            c, o = list(self._corr), list(self._ofs)
        if decimate > 1 and t:
            k = (len(t) - 1) % decimate
            t, v, c, o = t[k::decimate], v[k::decimate], c[k::decimate], o[k::decimate]
        return t, v, c, o

    def get_rates(self) -> tuple[int, int]:
        """(fresh publishes received, publishes sent) since the last call."""
        with self._lock:
            f, s = self._fresh_count, self._sent_count
            self._fresh_count = 0
            self._sent_count = 0
        return f, s

    def send_frame(self, data: bytes) -> None:
        """Queue a pre-built command frame for TX on the worker thread."""
        self._tx_queue.put(data)

    def start_csv(self, path: str) -> None:
        """Open the full-rate telemetry CSV (one row per fresh publish)."""
        with self._csv_lock:
            self._close_csv_locked()
            f = open(path, "w", newline="", encoding="utf-8")
            w = csv.writer(f)
            w.writerow(("pc_time", "t_mono") + TELEM_CSV_FIELDS + ("p_in_W", "p_out_W"))
            self._csv_file, self._csv_writer = f, w
            self._csv_last_flush = time.monotonic()

    def stop_csv(self) -> None:
        with self._csv_lock:
            self._close_csv_locked()

    def _close_csv_locked(self) -> None:
        if self._csv_file is not None:
            try:
                self._csv_file.close()
            except OSError:
                pass
        self._csv_file = None
        self._csv_writer = None

    def stop(self) -> None:
        """Stop the thread and close the CSV (blocks up to 2 s)."""
        self._stop_event.set()
        if self.is_alive():
            self.join(timeout=2.0)
        self.stop_csv()

    # ── Worker loop ──────────────────────────────────────────────────

    def run(self) -> None:
        ser = None
        parser = None
        while not self._stop_event.is_set():
            if ser is None:
                try:
                    ser = serial.Serial(
                        self.port, self.baudrate, timeout=0.05,
                        bytesize=serial.EIGHTBITS,
                        parity=serial.PARITY_ODD,
                        stopbits=serial.STOPBITS_ONE,
                    )
                    parser = BroadcastParser()
                    self.connected = True
                except (serial.SerialException, OSError):
                    self.connected = False
                    self._stop_event.wait(1.0)
                    continue
            try:
                while True:
                    try:
                        ser.write(self._tx_queue.get_nowait())
                    except queue.Empty:
                        break
                data = ser.read(4096)
                if data:
                    for b in parser.feed(data):
                        self._on_frame(b)
            except (serial.SerialException, OSError):
                self.connected = False
                try:
                    ser.close()
                except Exception:
                    pass
                ser = None
                self._stop_event.wait(1.0)
        if ser is not None:
            try:
                ser.close()
            except Exception:
                pass
        self.connected = False
        self.stop_csv()

    # NB: name must not collide with threading.Thread instance attributes
    # (Thread.__init__ sets self._handle on Python 3.13+).
    def _on_frame(self, b: BroadcastData) -> None:
        now = time.monotonic()
        with self._lock:
            # Flags/pong are valid even at counter == 0 (post-reset window);
            # telemetry/param fields are only valid once counter > 0.
            self.latest = b
            self.latest_wall = time.time()
            fresh = b.counter > 0 and b.counter != self._last_counter
            if fresh:
                if 0 < self._last_counter < b.counter:
                    self._sent_count += b.counter - self._last_counter
                else:
                    self._sent_count += 1   # first valid frame, or a reboot
                self._fresh_count += 1
                self._last_counter = b.counter
                self._t.append(now)
                self._vfly.append(b.vfly)
                self._corr.append(b.vfly_correction)
                self._ofs.append(b.vfly_ofs_applied)
                if b.param_id in PARAM_BY_ID:
                    self.params[b.param_id] = decode_param_bits(b.param_type, b.param_value)
        if fresh:
            self._csv_row(b, now)

    def _csv_row(self, b: BroadcastData, t_mono: float) -> None:
        with self._csv_lock:
            if self._csv_writer is None:
                return
            row = [f"{b.timestamp:.3f}", f"{t_mono:.3f}"]
            for name in TELEM_CSV_FIELDS:
                v = getattr(b, name)
                if name in _TELEM_HEX_FIELDS:
                    row.append(f"0x{v:08X}")
                elif isinstance(v, float):
                    row.append(f"{v:.6g}")
                else:
                    row.append(v)
            row.append(f"{b.power_in_W:.4f}")
            row.append(f"{b.power_out_W:.4f}")
            try:
                self._csv_writer.writerow(row)
                if t_mono - self._csv_last_flush >= self.CSV_FLUSH_S:
                    self._csv_file.flush()
                    self._csv_last_flush = t_mono
            except OSError:
                self._close_csv_locked()