From a5b391e7579b3dd45e5cc9d8c8447f3006210bff Mon Sep 17 00:00:00 2001
From: Brendan Haines <brendan.haines@gmail.com>
Date: Wed, 18 Dec 2024 23:31:34 -0700
Subject: [PATCH] untested: add some plot update stuff

---
 charon_vna/gui.py | 55 +++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 46 insertions(+), 9 deletions(-)

diff --git a/charon_vna/gui.py b/charon_vna/gui.py
index f154535..7e53652 100644
--- a/charon_vna/gui.py
+++ b/charon_vna/gui.py
@@ -1,13 +1,15 @@
 import sys
 from pathlib import Path
-from typing import List
+from typing import Callable, List, Tuple
 
 import matplotlib as mpl
 import numpy as np
+import xarray as xr
 from gui_helpers import FlowLayout
 from matplotlib import pyplot as plt
 from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg
 from matplotlib.ticker import EngFormatter
+from numpy import typing as npt
 from PySide6.QtCore import QSize, Qt
 from PySide6.QtGui import QAction, QKeySequence, QShortcut
 from PySide6.QtWidgets import (
@@ -21,6 +23,7 @@ from PySide6.QtWidgets import (
     QWidget,
 )
 from skrf import plotting as rf_plt
+from vna import Charon
 
 DEFAULT_CONFIG = dict(
     frequency=np.arange(1e9, 2e9, 11),  # Hz
@@ -29,9 +32,13 @@ DEFAULT_CONFIG = dict(
 
 
 class PlotWidget(QWidget):
+    enabled_ports: List[Tuple[int | str]]
+
     def __init__(self):
         super().__init__()
 
+        self.enabled_ports = [(1, 1)]
+
         layout = QVBoxLayout()
         self.setLayout(layout)
 
@@ -48,44 +55,74 @@ class PlotWidget(QWidget):
         # toolbar.addAction("blah")
         # self.addToolBar(toolbar)
 
-    def setup_rect(self):
+    def setup_rect(self) -> None:
         self.ax.grid(True)
         self.ax.xaxis.set_major_formatter(EngFormatter())
         self.ax.set_xlabel("Frequency [Hz]")
 
-    def setup_logmag(self, ylim: List[float] = [-30, 30]):
+    def update_rect(
+        self, data: xr.DataArray, func: Callable[[npt.ArrayLike[np.complex128]], npt.ArrayLike[float]]
+    ) -> None:
+        self.ax.set_xlim(data["frequency"].min().data, data["frequency"].max().data)
+        # remove old lines
+        for line in self.ax.lines:
+            line.remove()
+        for m, n in self.enabled_ports:
+            self.ax.plot(data["frequency"], func(data.sel(m=m, n=n)))
+
+    def setup_logmag(self, ylim: List[float] = [-30, 30]) -> None:
         self.setup_rect()
         self.ax.set_ylim(ylim)
         self.ax.set_ylabel("Amplitude [dB]")
 
-    def setup_phase(self):
+    def update_logmag(self, data: xr.DataArray) -> None:
+        self.update_rect(data, lambda s: 20 * np.log10(np.abs(s)))
+
+    def setup_phase(self) -> None:
         self.setup_rect()
         self.ax.set_ylim(-200, 200)
         self.ax.set_ylabel("Phase [deg]")
 
-    def setup_vswr(self):
+    def update_phase(self, data: xr.DataArray):
+        self.update_rect(data, lambda s: np.angle(s, deg=True))
+
+    def setup_vswr(self) -> None:
         self.setup_rect()
         self.ax.set_yticks(np.arange(1, 11))
         self.ax.set_ylim(1, 10)
         self.ax.set_ylabel("VSWR")
 
-    def setup_smith(self):
+    def update_vswr(self, data: xr.DataArray) -> None:
+        self.update_rect(data, lambda s: (1 + np.abs(s)) / (1 - np.abs(s)))
+
+    def setup_smith(self) -> None:
         self.ax.grid(False)
         self.ax.set_xlim(-1, 1)
         self.ax.set_ylim(-1, 1)
         self.ax.set_aspect("equal")
         rf_plt.smith(ax=self.ax, smithR=1, chart_type="z", draw_vswr=None)
 
+    def update_smith(self, data: xr.DataArray) -> None:
+        # remove old lines
+        for line in self.ax.lines:
+            line.remove()
+        for m, n in self.enabled_ports:
+            sel = data.sel(m=m, n=n)
+            self.ax.plot(sel.real, sel.imag)
+
 
 # Subclass QMainWindow to customize your application's main window
 class MainWindow(QMainWindow):
     config_path: Path | None
+    device: Charon
 
     def __init__(self):
         super().__init__()
 
         self.config_path = None
 
+        self.device = Charon("ip:192.168.3.1", frequency=DEFAULT_CONFIG["frequency"])
+
         mpl.use("QtAgg")
 
         self.setWindowTitle("Charon VNA")
@@ -139,21 +176,21 @@ class MainWindow(QMainWindow):
         widget.setLayout(window_layout)
         self.setCentralWidget(widget)
 
-    def saveas_config(self):
+    def saveas_config(self) -> None:
         print("Prompting for save path...")
         # TODO: prompt for config path
         self.config_path = Path(__file__).parent / "config.json"
         print(f"Config path is now {self.config_path.resolve()}")
         self.save_config()
 
-    def save_config(self):
+    def save_config(self) -> None:
         if self.config_path is None:
             self.saveas_config()
         else:
             print(f"saving config to {self.config_path.resolve()}")
             # TODO: save config
 
-    def load_config(self):
+    def load_config(self) -> None:
         print("loading config")
         # TODO: load config