radar_system/python_app/gui/controllers/renderers/trace.py

"""Magnitude and phase rendering with owned plots, curves and legends."""

from __future__ import annotations

from typing import Protocol

import numpy as np
import pyqtgraph as pg
from PyQt6.QtGui import QPen
from PyQt6.QtWidgets import QVBoxLayout, QWidget

from python_app.gui.controllers.renderers.common import CAPTURED_TRACE_PEN, OVERLAY_FILL, SERIES_COLORS, create_plot
from python_app.models.dataset_model import ComboKey, ResultCollection, TraceData
from python_app.models.gui_profile_schema import GuiPassThroughStateModel
from python_app.models.run_config_validation import parse_combos_from_text


class WarningLogger(Protocol):
    def __call__(self, text: str, *, details: str | None = None,
                 once_key: str | None = None) -> None: ...


CurveKey = tuple[ComboKey, int, str]
Series = tuple[CurveKey, np.ndarray, np.ndarray, QPen]


class _TracePanel:
    """One persistent curve surface and its combo legend."""

    def __init__(self, label: str, units: str) -> None:
        self.plot = create_plot("Pass-Through S21")
        item = self.plot.getPlotItem()
        assert item is not None
        self.item = item
        self.view = item.getViewBox()
        self.item.setLabel("left", label, units=units)
        self.item.setLabel("bottom", "Frequency", units="Hz")
        # Sweeps denser than the screen are drawn as per-bin min/max, so narrow peaks
        # and notches stay visible; ordinary sweeps keep every point.
        self.item.setDownsampling(auto=True, mode="peak")
        self.curves: dict[CurveKey, pg.PlotDataItem] = {}
        self.legend = self.item.addLegend(offset=(8, 8), brush=pg.mkBrush(*OVERLAY_FILL), colCount=2)
        self.legend.hide()
        self._legend_sources: dict[ComboKey, pg.PlotDataItem] = {}

    def clear(self) -> None:
        self.item.clear()
        self.curves.clear()
        self.sync_legend({})

    def update_curve(self, key: CurveKey, x: np.ndarray, y: np.ndarray, pen: QPen) -> pg.PlotDataItem:
        curve = self.curves.get(key)
        if curve is None:
            curve = pg.PlotDataItem(antialias=False)
            self.curves[key] = curve
            self.item.addItem(curve)
        curve.setData(x, y, pen=pen)  # one update; a separate setPen would redraw the old data first
        return curve

    def retain_curves(self, active_keys: set[CurveKey]) -> None:
        for key in self.curves.keys() - active_keys:
            self.item.removeItem(self.curves.pop(key))

    def sync_legend(self, sources: dict[ComboKey, pg.PlotDataItem]) -> None:
        if self._legend_sources == sources:
            return
        self.legend.clear()
        for combo in sorted(sources, key=lambda key: (key.input, key.output)):
            self.legend.addItem(sources[combo], f"in{combo.input}/out{combo.output}")
        self.legend.setVisible(bool(sources))
        self._legend_sources = sources


class TraceRenderer:
    """Render result traces or a captured trace using explicit display settings."""

    def __init__(self, log_warning: WarningLogger) -> None:
        self._log_warning = log_warning
        self.widget = QWidget()
        layout = QVBoxLayout(self.widget)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(6)
        self.magnitude = _TracePanel("Magnitude", "dB")
        self.phase = _TracePanel("Phase", "deg")
        self.magnitude.item.showAxis("bottom", show=False)
        for panel in (self.magnitude, self.phase):
            layout.addWidget(panel.plot, stretch=1)

    def clear(self) -> None:
        self.magnitude.clear()
        self.phase.clear()

    def draw(self, collection: ResultCollection, settings: GuiPassThroughStateModel) -> bool:
        combo_filter = self._combo_filter(settings.combo_filter)
        series: list[Series] = []
        # The color follows the pair's place in the run, so filtering never repaints shown pairs.
        for position, block in enumerate(collection.blocks):
            if combo_filter is not None and block.combo not in combo_filter:
                continue
            pen = pg.mkPen(SERIES_COLORS[position % len(SERIES_COLORS)], width=2)
            for index, payload in enumerate(block.payloads):
                # Another mode's result (a B-scan range profile) is not S21 over frequency.
                if payload.processing_name != "pass_through" or payload.kind != 1:
                    continue
                series.append(((block.combo, index, payload.processing_name),
                               payload.frequency_hz, payload.trace, pen))
        return self._render(series, settings, "Pass-Through S21")

    def draw_trace(self, trace: TraceData, title: str, settings: GuiPassThroughStateModel) -> None:
        self.clear()
        self._render([((trace.combo, 0, "__single_trace__"), trace.frequency_hz,
                       trace.s21, CAPTURED_TRACE_PEN)], settings, title, single=True)

    def _combo_filter(self, text: str) -> set[ComboKey] | None:
        text = text.strip()
        if not text:
            return None
        try:
            return {ComboKey(combo.input, combo.output) for combo in parse_combos_from_text(text)}
        except Exception as exc:  # noqa: BLE001
            self._log_warning("Invalid pass-through switch-combo filter.",
                              details=f"{exc}\nExpected format: input:output,input:output",
                              once_key=f"pass_through_combo_filter_invalid_{text}")
            return None

    def _render(self, series: list[Series], settings: GuiPassThroughStateModel,
                title: str, *, single: bool = False) -> bool:
        mag, phase = self.magnitude, self.phase
        mag.plot.setVisible(settings.show_magnitude)
        phase.plot.setVisible(settings.show_phase)
        if not settings.show_magnitude and not settings.show_phase:
            self.clear()
            return False
        mag.item.showAxis("bottom", show=not settings.show_phase)
        active: set[CurveKey] = set()
        mag_sources: dict[ComboKey, pg.PlotDataItem] = {}
        phase_sources: dict[ComboKey, pg.PlotDataItem] = {}
        x_min, x_max = np.inf, -np.inf
        for key, frequency, samples, pen in series:
            if not samples.size or frequency.size != samples.size:
                continue
            active.add(key)
            x_min, x_max = min(x_min, float(np.min(frequency))), max(x_max, float(np.max(frequency)))
            if settings.show_magnitude:
                magnitude = 20.0 * np.log10(np.maximum(np.abs(samples), 1e-12))
                curve = mag.update_curve(key, frequency, magnitude, pen)
                mag_sources.setdefault(key[0], curve)
            if settings.show_phase:
                angles = np.angle(samples)
                if settings.unwrap_phase:
                    angles = np.unwrap(angles)
                curve = phase.update_curve(key, frequency, np.degrees(angles), pen)
                phase_sources.setdefault(key[0], curve)

        for panel, visible, sources in ((mag, settings.show_magnitude, mag_sources),
                                        (phase, settings.show_phase, phase_sources)):
            panel.retain_curves(active if visible else set())
            panel.sync_legend(sources if not single else {})
            if visible:
                panel.item.setTitle(title)
                if active and np.isfinite(x_min) and np.isfinite(x_max):
                    panel.view.setXRange(x_min, x_max, padding=0.02)
        if settings.show_magnitude:
            mag.view.enableAutoRange(y=not settings.fixed_y_enabled)
            if settings.fixed_y_enabled:
                mag.view.setYRange(min(settings.y_min_db, settings.y_max_db),
                                   max(settings.y_min_db, settings.y_max_db), padding=0.0)
        if settings.show_phase:
            phase.view.enableAutoRange(y=settings.unwrap_phase)
            if not settings.unwrap_phase:
                phase.view.setYRange(-180.0, 180.0, padding=0.02)
        return bool(active)