From 5fb7fbc8ec711bf2b97443e484736dea17348174 Mon Sep 17 00:00:00 2001
From: Simon Renblad <srenblad@m-labs.hk>
Date: Tue, 16 Apr 2024 14:26:03 +0800
Subject: [PATCH] moninj: _AddChannelDialog

---
 artiq/dashboard/moninj.py | 71 ++++++++++++++++++++++++++++++++++++++-
 1 file changed, 70 insertions(+), 1 deletion(-)

diff --git a/artiq/dashboard/moninj.py b/artiq/dashboard/moninj.py
index 77ca07a74..24ca3e43d 100644
--- a/artiq/dashboard/moninj.py
+++ b/artiq/dashboard/moninj.py
@@ -9,6 +9,7 @@ from artiq.coredevice.comm_moninj import CommMonInj, TTLOverride, TTLProbe
 from artiq.coredevice.ad9912_reg import AD9912_SER_CONF
 from artiq.gui.tools import LayoutWidget
 from artiq.gui.flowlayout import FlowLayout
+from artiq.gui.models import DictSyncTreeSepModel
 
 
 logger = logging.getLogger(__name__)
@@ -410,6 +411,58 @@ class _DACHandler:
         return (self.spi_channel, self.channel)
 
 
+class Model(DictSyncTreeSepModel):
+    def __init__(self, init):
+        DictSyncTreeSepModel.__init__(self, "/", ["Channels"], init)
+
+    def clear(self):
+        for k in self.backing_store:
+            self._del_item(self, k.split(self.separator))
+        self.backing_store.clear()
+
+    def update(self, d):
+        for k, v in d.items():
+            self[str(k)] = v
+
+
+class _AddChannelDialog(QtWidgets.QDialog):
+    def __init__(self, parent, model):
+        QtWidgets.QDialog.__init__(self, parent=parent)
+        self.setContextMenuPolicy(QtCore.Qt.ActionsContextMenu)
+        self.setWindowTitle("Add channels")
+
+        layout = QtWidgets.QVBoxLayout()
+        self.setLayout(layout)
+
+        self._model = model
+        self._tree_view = QtWidgets.QTreeView()
+        self._tree_view.setHeaderHidden(True)
+        self._tree_view.setSelectionBehavior(
+            QtWidgets.QAbstractItemView.SelectItems)
+        self._tree_view.setSelectionMode(
+            QtWidgets.QAbstractItemView.ExtendedSelection)
+        self._tree_view.setModel(self._model)
+        layout.addWidget(self._tree_view)
+
+        self._button_box = QtWidgets.QDialogButtonBox(
+            QtWidgets.QDialogButtonBox.Ok | QtWidgets.QDialogButtonBox.Cancel
+        )
+        self._button_box.setCenterButtons(True)
+        self._button_box.accepted.connect(self.add_channels)
+        self._button_box.rejected.connect(self.reject)
+        layout.addWidget(self._button_box)
+
+    def add_channels(self):
+        selection = self._tree_view.selectedIndexes()
+        channels = []
+        for select in selection:
+            key = self._model.index_to_key(select)
+            if key is not None:
+                channels.append(self._model[key].ref)
+        self.channels = channels
+        self.accept()
+
+
 _HandlerDesc = namedtuple("_HandlerDesc", "uid comment cls arguments")
 
 
@@ -493,6 +546,7 @@ class _DeviceManager:
         self.dds_handlers = dict()
         self.dac_cb = lambda: None
         self.dac_handlers = dict()
+        self.channels_cb = lambda: None
 
     def init_ddb(self, ddb):
         self.ddb = ddb
@@ -547,6 +601,7 @@ class _DeviceManager:
             self.ttl_cb()
             self.dds_cb()
             self.dac_cb()
+            self.channels_cb()
 
         self.description = description
 
@@ -794,8 +849,17 @@ class _MonInjDock(QtWidgets.QDockWidget):
                          QtWidgets.QDockWidget.DockWidgetFloatable)
         layout = LayoutWidget()
         self.setWidget(layout)
+        self._channel_model = Model({})
+        self.add_channel_dialog = _AddChannelDialog(self, self._channel_model)
+        add_channel_btn = QtWidgets.QToolButton()
+        add_channel_btn.setToolTip("Add channels...")
+        add_channel_btn.setIcon(
+            QtWidgets.QApplication.style().standardIcon(
+                QtWidgets.QStyle.SP_FileDialogListView))
+        add_channel_btn.clicked.connect(self.add_channel_dialog.open)
+        layout.addWidget(add_channel_btn, 0, 0, colspan=1)
         scroll_area = QtWidgets.QScrollArea()
-        layout.addWidget(scroll_area)
+        layout.addWidget(scroll_area, 1, 0, colspan=10)
         self.grid = FlowLayout()
         grid_widget = QtWidgets.QWidget()
         grid_widget.setLayout(self.grid)
@@ -806,6 +870,10 @@ class _MonInjDock(QtWidgets.QDockWidget):
         for handler in sorted(handlers, key=lambda h: h.sort_key()):
             self.grid.addWidget(handler.widget)
 
+    def set_channels(self, handlers):
+        self._channel_model.clear()
+        self._channel_model.update(handlers)
+
 
 class MonInj:
     def __init__(self, schedule_ctl):
@@ -817,6 +885,7 @@ class MonInj:
         self.dm.ttl_cb = lambda: self.ttl_dock.layout_widgets(self.dm.ttl_handlers.values())
         self.dm.dds_cb = lambda: self.dds_dock.layout_widgets(self.dm.dds_handlers.values())
         self.dm.dac_cb = lambda: self.dac_dock.layout_widgets(self.dm.dac_handlers.values())
+        self.dm.channels_cb = lambda: self.ttl_dock.set_channels(self.dm.handlers_by_uid)
 
     async def stop(self):
         if self.dm is not None: