From 33495cfe03b363f18db61d8af2983f49027b7a43 Mon Sep 17 00:00:00 2001
From: wyzula-jan <jan.wyzula@psi.ch>
Date: Sun, 14 Jul 2024 23:19:40 +0200
Subject: [PATCH] fix(waveform): colormaps of curves can be changed and
 normalised

feat(waveform): colormap can be changed from curve dialog

fix(curve_dialog): default dialog parameters fixed

curve Dialog colormap WIP
---
 .../jupyter_console/jupyter_console_window.py |  13 +-
 .../widgets/figure/plots/waveform/waveform.py |  69 +--
 .../curve_dialog/curve_dialog.py              |  18 +-
 .../curve_dialog/curve_dialog.ui              | 441 +++++++++---------
 .../widgets/waveform/waveform_widget.py       |   9 +
 5 files changed, 291 insertions(+), 259 deletions(-)

diff --git a/bec_widgets/examples/jupyter_console/jupyter_console_window.py b/bec_widgets/examples/jupyter_console/jupyter_console_window.py
index 821ac50f..ed3ec96f 100644
--- a/bec_widgets/examples/jupyter_console/jupyter_console_window.py
+++ b/bec_widgets/examples/jupyter_console/jupyter_console_window.py
@@ -49,8 +49,9 @@ class JupyterConsoleWindow(QWidget):  # pragma: no cover:
                     "d0": self.d0,
                     "d1": self.d1,
                     "d2": self.d2,
-                    "plt": self.plt,
+                    "wave": self.wave,
                     "bar": self.bar,
+                    "cm": self.colormap,
                 }
             )
 
@@ -165,12 +166,16 @@ class JupyterConsoleWindow(QWidget):  # pragma: no cover:
         self.fig1.plot(x_name="samx", y_name="bpm3a")
 
         self.d2 = self.dock.add_dock(name="dock_2", position="bottom")
-        self.fig2 = self.d2.add_widget("BECFigure", row=0, col=0)
-        self.plt = self.fig2.plot(x_name="samx", y_name="bpm3a")
-        self.plt.plot(x_name="samx", y_name="bpm4i", dap="GaussianModel")
+        self.wave = self.d2.add_widget("BECWaveformWidget", row=0, col=0)
+        # self.wave.plot(x_name="samx", y_name="bpm3a")
+        # self.wave.plot(x_name="samx", y_name="bpm4i", dap="GaussianModel")
         self.bar = self.d2.add_widget("RingProgressBar", row=0, col=1)
         self.bar.set_diameter(200)
 
+        self.d3 = self.dock.add_dock(name="dock_3", position="bottom")
+        self.colormap = pg.GradientWidget()
+        self.d3.add_widget(self.colormap, row=0, col=0)
+
         self.dock.save_state()
 
     def closeEvent(self, event):
diff --git a/bec_widgets/widgets/figure/plots/waveform/waveform.py b/bec_widgets/widgets/figure/plots/waveform/waveform.py
index e609889a..78ef1fe4 100644
--- a/bec_widgets/widgets/figure/plots/waveform/waveform.py
+++ b/bec_widgets/widgets/figure/plots/waveform/waveform.py
@@ -10,7 +10,7 @@ import pyqtgraph as pg
 from bec_lib import messages
 from bec_lib.device import ReadoutPriority
 from bec_lib.endpoints import MessageEndpoints
-from pydantic import Field, ValidationError
+from pydantic import Field, ValidationError, field_validator
 from qtpy.QtCore import Signal as pyqtSignal
 from qtpy.QtCore import Slot as pyqtSlot
 from qtpy.QtWidgets import QWidget
@@ -26,13 +26,16 @@ from bec_widgets.widgets.figure.plots.waveform.waveform_curve import (
 
 
 class Waveform1DConfig(SubplotConfig):
-    color_palette: Literal["plasma", "viridis", "inferno", "magma"] = Field(
-        "plasma", description="The color palette of the figure widget."
+    color_palette: Optional[str] = Field(
+        "plasma", description="The color palette of the figure widget.", validate_default=True
     )
     curves: dict[str, CurveConfig] = Field(
         {}, description="The list of curves to be added to the 1D waveform widget."
     )
 
+    model_config: dict = {"validate_assignment": True}
+    _validate_color_map_z = field_validator("color_palette")(Colors.validate_color_map)
+
 
 class BECWaveform(BECPlotBase):
     READOUT_PRIORITY_HANDLER = {
@@ -63,6 +66,7 @@ class BECWaveform(BECPlotBase):
         "set_x_lim",
         "set_y_lim",
         "set_grid",
+        "set_colormap",
         "lock_aspect_ratio",
         "remove",
         "clear_all",
@@ -954,6 +958,22 @@ class BECWaveform(BECPlotBase):
         current_label = "" if self.config.axis.x_label is None else self.config.axis.x_label
         self.plot_item.setLabel("bottom", f"{current_label}{self._x_axis_mode['label_suffix']}")
 
+    def set_colormap(self, colormap: str | None = None):
+        """
+        Set the colormap of the plot widget.
+
+        Args:
+            colormap(str, optional): Scale the colors of curves to colormap. If None, use the default color palette.
+        """
+        if colormap is not None:
+            self.config.color_palette = colormap
+
+        colors = Colors.golden_angle_color(
+            colormap=self.config.color_palette, num=len(self.plot_item.curves) + 1, format="HEX"
+        )
+        for curve, color in zip(self.curves, colors):
+            curve.set_color(color)
+
     def setup_dap(self, old_scan_id: str | None, new_scan_id: str | None):
         """
         Setup DAP for the new scan.
@@ -1214,49 +1234,6 @@ class BECWaveform(BECPlotBase):
                 x_data = []
         return x_data
 
-    # def _get_x_data(self, curve: BECCurve, y_name: str, y_entry: str) -> list | np.ndarray | None:
-    #     """
-    #     Get the x data for the curve with the decision logic based on the curve configuration:
-    #         - If x is called 'timestamp', use the timestamp data from the scan item.
-    #         - If x is called 'index', use the rolling index.
-    #         - If x is a custom signal, use the data from the scan item.
-    #         - If x is not specified, use the first device from the scan report.
-    #
-    #     Args:
-    #         curve(BECCurve): The curve object.
-    #
-    #     Returns:
-    #         list|np.ndarray|None: X data for the curve.
-    #     """
-    #     x_data = None
-    #     if curve.config.signals.x is not None:
-    #         if curve.config.signals.x.name == "timestamp":
-    #             timestamps = self.scan_item.data[y_name][y_entry].timestamps
-    #             x_data = self.convert_timestamps(timestamps)
-    #         elif curve.config.signals.x.name == "index":
-    #             x_data = None
-    #         else:
-    #             x_name = curve.config.signals.x.name
-    #             x_entry = curve.config.signals.x.entry
-    #             try:
-    #                 x_data = self.scan_item.data[x_name][x_entry].val
-    #             except TypeError:
-    #                 x_data = []
-    #     else:
-    #         if len(self._curves_data["async"]) > 0:
-    #             x_data = None
-    #         else:
-    #             x_name = self.scan_item.status_message.info["scan_report_devices"][0]
-    #             x_entry = self.entry_validator.validate_signal(x_name, None)
-    #             x_data = self.scan_item.data[x_name][x_entry].val
-    #             self._x_axis_mode["label_suffix"] = f" [auto: {x_name}-{x_entry}]"
-    #             current_label = "" if self.config.axis.x_label is None else self.config.axis.x_label
-    #             self.plot_item.setLabel(
-    #                 "bottom", f"{current_label}{self._x_axis_mode['label_suffix']}"
-    #             )
-    #
-    #     return x_data
-
     def _make_z_gradient(self, data_z: list | np.ndarray, colormap: str) -> list | None:
         """
         Make a gradient color for the z values.
diff --git a/bec_widgets/widgets/waveform/waveform_toolbar/curve_dialog/curve_dialog.py b/bec_widgets/widgets/waveform/waveform_toolbar/curve_dialog/curve_dialog.py
index 71990343..aa85dc25 100644
--- a/bec_widgets/widgets/waveform/waveform_toolbar/curve_dialog/curve_dialog.py
+++ b/bec_widgets/widgets/waveform/waveform_toolbar/curve_dialog/curve_dialog.py
@@ -10,7 +10,7 @@ from qtpy.QtCore import Slot
 from qtpy.QtWidgets import QVBoxLayout
 
 from bec_widgets.qt_utils.settings_dialog import SettingWidget
-from bec_widgets.utils import UILoader
+from bec_widgets.utils import UILoader, Colors
 from bec_widgets.widgets.color_button.color_button import ColorButton
 from bec_widgets.widgets.device_line_edit.device_line_edit import DeviceLineEdit
 from bec_widgets.widgets.figure.plots.plot_base import AxisConfig
@@ -30,6 +30,7 @@ class CurveSettings(SettingWidget):
 
         self.ui.add_curve.clicked.connect(self.add_curve)
         self.ui.x_mode.currentIndexChanged.connect(self.set_x_mode)
+        self.ui.normalize_colors.clicked.connect(self.change_colormap)
 
     @Slot(dict)
     def display_current_settings(self, config: dict | BaseModel):
@@ -39,6 +40,8 @@ class CurveSettings(SettingWidget):
         x_name = self.target_widget.waveform._x_axis_mode["name"]
         x_entry = self.target_widget.waveform._x_axis_mode["entry"]
         self._setup_x_box(x_name, x_entry)
+        cm = self.target_widget.config.color_palette
+        self.ui.color_map_selector.combo.setCurrentText(cm)
 
         for label, curve in curves.items():
             row_count = self.ui.scan_table.rowCount()
@@ -65,6 +68,15 @@ class CurveSettings(SettingWidget):
             self.ui.x_name.setEnabled(True)
             self.ui.x_entry.setEnabled(True)
 
+    @Slot()
+    def change_colormap(self):
+        cm = self.ui.color_map_selector.combo.currentText()
+        rows = self.ui.scan_table.rowCount()
+        colors = Colors.golden_angle_color(colormap=cm, num=rows + 1, format="HEX")
+        for row, color in zip(range(rows), colors):
+            self.ui.scan_table.cellWidget(row, 2).setColor(color)
+        self.target_widget.set_colormap(cm)
+
     @Slot()
     def accept_changes(self):
         self.accept_scan_curve_changes()
@@ -131,15 +143,19 @@ class ScanRow(QObject):
         self.entry_line_edit = QLineEdit()
 
         # Styling
+        default_color = Colors.golden_angle_color(colormap="magma", num=row + 1, format="HEX")[-1]
         self.color_button = ColorButton()
+        self.color_button.setColor(default_color)
         self.style_combo = StyleComboBox()
         self.width = QSpinBox()
         self.width.setMinimum(1)
         self.width.setMaximum(20)
+        self.width.setValue(2)
 
         self.symbol_size = QSpinBox()
         self.symbol_size.setMinimum(1)
         self.symbol_size.setMaximum(20)
+        self.symbol_size.setValue(5)
 
         self.table_widget = table_widget
         self.row = row
diff --git a/bec_widgets/widgets/waveform/waveform_toolbar/curve_dialog/curve_dialog.ui b/bec_widgets/widgets/waveform/waveform_toolbar/curve_dialog/curve_dialog.ui
index 437169d5..094595e1 100644
--- a/bec_widgets/widgets/waveform/waveform_toolbar/curve_dialog/curve_dialog.ui
+++ b/bec_widgets/widgets/waveform/waveform_toolbar/curve_dialog/curve_dialog.ui
@@ -6,92 +6,92 @@
    <rect>
     <x>0</x>
     <y>0</y>
-       <width>720</width>
-       <height>806</height>
+    <width>720</width>
+    <height>806</height>
    </rect>
   </property>
   <property name="windowTitle">
    <string>Form</string>
   </property>
   <layout class="QVBoxLayout" name="verticalLayout_2">
-      <property name="leftMargin">
-          <number>2</number>
-      </property>
-      <property name="topMargin">
-          <number>2</number>
-      </property>
-      <property name="rightMargin">
-          <number>2</number>
-      </property>
-      <property name="bottomMargin">
-          <number>2</number>
-      </property>
+   <property name="leftMargin">
+    <number>2</number>
+   </property>
+   <property name="topMargin">
+    <number>2</number>
+   </property>
+   <property name="rightMargin">
+    <number>2</number>
+   </property>
+   <property name="bottomMargin">
+    <number>2</number>
+   </property>
    <item>
     <widget class="QGroupBox" name="x_group_box">
      <property name="title">
       <string>X Axis</string>
      </property>
-        <layout class="QHBoxLayout" name="horizontalLayout" stretch="0,0,0,0,1,3,0,1,3">
-            <item>
-                <widget class="QLabel" name="x_mode_label">
-                    <property name="text">
-                        <string>X Axis Mode</string>
-                    </property>
-                </widget>
-            </item>
-            <item>
-                <widget class="QComboBox" name="x_mode">
-                    <item>
-                        <property name="text">
-                            <string>best_effort</string>
-                        </property>
-                    </item>
-                    <item>
-                        <property name="text">
-                            <string>device</string>
-                        </property>
-                    </item>
-                    <item>
-                        <property name="text">
-                            <string>index</string>
-                        </property>
-                    </item>
-                    <item>
-                        <property name="text">
-                            <string>timestamp</string>
-                        </property>
-                    </item>
-                </widget>
-            </item>
-            <item>
-                <spacer name="horizontalSpacer_3">
-                    <property name="orientation">
-                        <enum>Qt::Orientation::Horizontal</enum>
-                    </property>
-                    <property name="sizeHint" stdset="0">
-                        <size>
-                            <width>40</width>
-                            <height>20</height>
-                        </size>
-                    </property>
-                </spacer>
-            </item>
-            <item>
-                <widget class="Line" name="line_2">
-                    <property name="frameShadow">
-                        <enum>QFrame::Shadow::Sunken</enum>
-                    </property>
-                    <property name="lineWidth">
-                        <number>1</number>
-                    </property>
-                    <property name="midLineWidth">
-                        <number>0</number>
-                    </property>
-                    <property name="orientation">
-                        <enum>Qt::Orientation::Vertical</enum>
-                    </property>
-                </widget>
-            </item>
+     <layout class="QHBoxLayout" name="horizontalLayout" stretch="0,0,0,0,1,3,0,1,3">
+      <item>
+       <widget class="QLabel" name="x_mode_label">
+        <property name="text">
+         <string>X Axis Mode</string>
+        </property>
+       </widget>
+      </item>
+      <item>
+       <widget class="QComboBox" name="x_mode">
+        <item>
+         <property name="text">
+          <string>best_effort</string>
+         </property>
+        </item>
+        <item>
+         <property name="text">
+          <string>device</string>
+         </property>
+        </item>
+        <item>
+         <property name="text">
+          <string>index</string>
+         </property>
+        </item>
+        <item>
+         <property name="text">
+          <string>timestamp</string>
+         </property>
+        </item>
+       </widget>
+      </item>
+      <item>
+       <spacer name="horizontalSpacer_3">
+        <property name="orientation">
+         <enum>Qt::Orientation::Horizontal</enum>
+        </property>
+        <property name="sizeHint" stdset="0">
+         <size>
+          <width>40</width>
+          <height>20</height>
+         </size>
+        </property>
+       </spacer>
+      </item>
+      <item>
+       <widget class="Line" name="line_2">
+        <property name="frameShadow">
+         <enum>QFrame::Shadow::Sunken</enum>
+        </property>
+        <property name="lineWidth">
+         <number>1</number>
+        </property>
+        <property name="midLineWidth">
+         <number>0</number>
+        </property>
+        <property name="orientation">
+         <enum>Qt::Orientation::Vertical</enum>
+        </property>
+       </widget>
+      </item>
       <item>
        <widget class="QLabel" name="x_name_label">
         <property name="text">
@@ -157,6 +157,78 @@
             </property>
            </widget>
           </item>
+          <item row="0" column="3">
+           <widget class="ColormapSelector" name="color_map_selector">
+            <property name="minimumSize">
+             <size>
+              <width>150</width>
+              <height>0</height>
+             </size>
+            </property>
+            <property name="colormaps" stdset="0">
+             <stringlist>
+              <string>inferno</string>
+              <string>viridis</string>
+              <string>plasma</string>
+              <string>magma</string>
+             </stringlist>
+            </property>
+           </widget>
+          </item>
+          <item row="1" column="0" colspan="6">
+           <widget class="QTableWidget" name="scan_table">
+            <property name="rowCount">
+             <number>0</number>
+            </property>
+            <attribute name="horizontalHeaderCascadingSectionResizes">
+             <bool>false</bool>
+            </attribute>
+            <attribute name="horizontalHeaderShowSortIndicator" stdset="0">
+             <bool>true</bool>
+            </attribute>
+            <attribute name="horizontalHeaderStretchLastSection">
+             <bool>true</bool>
+            </attribute>
+            <attribute name="verticalHeaderStretchLastSection">
+             <bool>false</bool>
+            </attribute>
+            <column>
+             <property name="text">
+              <string>Name</string>
+             </property>
+            </column>
+            <column>
+             <property name="text">
+              <string>Entry</string>
+             </property>
+            </column>
+            <column>
+             <property name="text">
+              <string>Color</string>
+             </property>
+            </column>
+            <column>
+             <property name="text">
+              <string>Style</string>
+             </property>
+            </column>
+            <column>
+             <property name="text">
+              <string>Width</string>
+             </property>
+            </column>
+            <column>
+             <property name="text">
+              <string>Symbol Size</string>
+             </property>
+            </column>
+            <column>
+             <property name="text">
+              <string>Delete</string>
+             </property>
+            </column>
+           </widget>
+          </item>
           <item row="0" column="1">
            <spacer name="horizontalSpacer">
             <property name="orientation">
@@ -170,23 +242,60 @@
             </property>
            </spacer>
           </item>
+          <item row="0" column="2">
+           <widget class="QPushButton" name="normalize_colors">
+            <property name="text">
+             <string>Normalize Colors</string>
+            </property>
+           </widget>
+          </item>
+         </layout>
+        </widget>
+        <widget class="QWidget" name="tab_DAP">
+         <attribute name="title">
+          <string>DAP</string>
+         </attribute>
+         <layout class="QGridLayout" name="gridLayout_2">
+          <property name="leftMargin">
+           <number>5</number>
+          </property>
+          <property name="topMargin">
+           <number>5</number>
+          </property>
+          <property name="rightMargin">
+           <number>5</number>
+          </property>
+          <property name="bottomMargin">
+           <number>5</number>
+          </property>
+          <item row="0" column="0">
+           <widget class="QPushButton" name="add_dap">
+            <property name="text">
+             <string>Add DAP</string>
+            </property>
+           </widget>
+          </item>
+          <item row="0" column="1">
+           <spacer name="horizontalSpacer_2">
+            <property name="orientation">
+             <enum>Qt::Orientation::Horizontal</enum>
+            </property>
+            <property name="sizeHint" stdset="0">
+             <size>
+              <width>585</width>
+              <height>20</height>
+             </size>
+            </property>
+           </spacer>
+          </item>
           <item row="1" column="0" colspan="2">
-           <widget class="QTableWidget" name="scan_table">
-               <property name="rowCount">
-                   <number>0</number>
-               </property>
-               <attribute name="horizontalHeaderCascadingSectionResizes">
-                   <bool>false</bool>
-               </attribute>
-               <attribute name="horizontalHeaderShowSortIndicator" stdset="0">
-                   <bool>true</bool>
-               </attribute>
-               <attribute name="horizontalHeaderStretchLastSection">
-                   <bool>true</bool>
-               </attribute>
-               <attribute name="verticalHeaderStretchLastSection">
-                   <bool>false</bool>
-               </attribute>
+           <widget class="QTableWidget" name="dap_table">
+            <attribute name="horizontalHeaderShowSortIndicator" stdset="0">
+             <bool>true</bool>
+            </attribute>
+            <attribute name="horizontalHeaderStretchLastSection">
+             <bool>true</bool>
+            </attribute>
             <column>
              <property name="text">
               <string>Name</string>
@@ -197,129 +306,40 @@
               <string>Entry</string>
              </property>
             </column>
-               <column>
-                   <property name="text">
-                       <string>Color</string>
-                   </property>
-               </column>
-               <column>
-                   <property name="text">
-                       <string>Style</string>
-                   </property>
-               </column>
-               <column>
-                   <property name="text">
-                       <string>Width</string>
-                   </property>
-               </column>
-               <column>
-                   <property name="text">
-                       <string>Symbol Size</string>
-                   </property>
-               </column>
-               <column>
-                   <property name="text">
-                       <string>Delete</string>
+            <column>
+             <property name="text">
+              <string>Model</string>
+             </property>
+            </column>
+            <column>
+             <property name="text">
+              <string>Color</string>
+             </property>
+            </column>
+            <column>
+             <property name="text">
+              <string>Style</string>
+             </property>
+            </column>
+            <column>
+             <property name="text">
+              <string>Width</string>
+             </property>
+            </column>
+            <column>
+             <property name="text">
+              <string>Symbol Size</string>
+             </property>
+            </column>
+            <column>
+             <property name="text">
+              <string>Delete</string>
              </property>
             </column>
            </widget>
           </item>
          </layout>
         </widget>
-        <widget class="QWidget" name="tab_DAP">
-         <attribute name="title">
-          <string>DAP</string>
-         </attribute>
-            <layout class="QGridLayout" name="gridLayout_2">
-                <property name="leftMargin">
-                    <number>5</number>
-                </property>
-                <property name="topMargin">
-                    <number>5</number>
-                </property>
-                <property name="rightMargin">
-                    <number>5</number>
-                </property>
-                <property name="bottomMargin">
-                    <number>5</number>
-                </property>
-                <item row="0" column="0">
-                    <widget class="QPushButton" name="add_dap">
-                        <property name="text">
-                            <string>Add DAP</string>
-                        </property>
-                    </widget>
-                </item>
-                <item row="0" column="1">
-                    <spacer name="horizontalSpacer_2">
-                        <property name="orientation">
-                            <enum>Qt::Orientation::Horizontal</enum>
-                        </property>
-                        <property name="sizeHint" stdset="0">
-                            <size>
-                                <width>585</width>
-                                <height>20</height>
-                            </size>
-                        </property>
-                    </spacer>
-                </item>
-                <item row="1" column="0" colspan="2">
-                    <widget class="QTableWidget" name="dap_table">
-                        <attribute name="horizontalHeaderShowSortIndicator" stdset="0">
-                            <bool>true</bool>
-                        </attribute>
-                        <attribute name="horizontalHeaderStretchLastSection">
-                            <bool>true</bool>
-                        </attribute>
-                        <column>
-                            <property name="text">
-                                <string>Name</string>
-                            </property>
-                        </column>
-                        <column>
-                            <property name="text">
-                                <string>Entry</string>
-                            </property>
-                        </column>
-                        <column>
-                            <property name="text">
-                                <string>Model</string>
-                            </property>
-                        </column>
-                        <column>
-                            <property name="text">
-                                <string>Color</string>
-                            </property>
-                        </column>
-                        <column>
-                            <property name="text">
-                                <string>Style</string>
-                            </property>
-                        </column>
-                        <column>
-                            <property name="text">
-                                <string>Width</string>
-                            </property>
-                        </column>
-                        <column>
-                            <property name="text">
-                                <string>Symbol Size</string>
-                            </property>
-                        </column>
-                        <column>
-                            <property name="text">
-                                <string>Delete</string>
-                            </property>
-                        </column>
-                    </widget>
-                </item>
-            </layout>
-        </widget>
-        <widget class="QWidget" name="tab_custom">
-         <attribute name="title">
-          <string>Custom</string>
-         </attribute>
-        </widget>
        </widget>
       </item>
      </layout>
@@ -333,6 +353,11 @@
    <extends>QLineEdit</extends>
    <header>device_line_edit</header>
   </customwidget>
+  <customwidget>
+   <class>ColormapSelector</class>
+   <extends>QWidget</extends>
+   <header>color_map_selector</header>
+  </customwidget>
  </customwidgets>
  <resources/>
  <connections/>
diff --git a/bec_widgets/widgets/waveform/waveform_widget.py b/bec_widgets/widgets/waveform/waveform_widget.py
index dac4b99c..2f412d14 100644
--- a/bec_widgets/widgets/waveform/waveform_widget.py
+++ b/bec_widgets/widgets/waveform/waveform_widget.py
@@ -175,6 +175,15 @@ class BECWaveformWidget(BECConnector, QWidget):
         """
         return self.waveform.get_curve(identifier)
 
+    def set_colormap(self, colormap: str):
+        """
+        Set the colormap of the plot widget.
+
+        Args:
+            colormap(str, optional): Scale the colors of curves to colormap. If None, use the default color palette.
+        """
+        self.waveform.set_colormap(colormap)
+
     def set_x(self, x_name: str, x_entry: str | None = None):
         """
         Change the x axis of the plot widget.