From 8e7885f36dd2812e3285c4d2d101212055644c7b Mon Sep 17 00:00:00 2001
From: wyzula-jan <133381102+wyzula-jan@users.noreply.github.com>
Date: Mon, 4 Sep 2023 10:54:59 +0200
Subject: [PATCH] fix: user selected colors are preserved with the new scan

---
 .../{config.yaml => config_example.yaml}      | 30 +++++++++++++-
 bec_widgets/examples/extreme/extreme.py       | 40 +++++++++++++++----
 2 files changed, 62 insertions(+), 8 deletions(-)
 rename bec_widgets/examples/extreme/{config.yaml => config_example.yaml} (64%)

diff --git a/bec_widgets/examples/extreme/config.yaml b/bec_widgets/examples/extreme/config_example.yaml
similarity index 64%
rename from bec_widgets/examples/extreme/config.yaml
rename to bec_widgets/examples/extreme/config_example.yaml
index 4f804a56..845ae99b 100644
--- a/bec_widgets/examples/extreme/config.yaml
+++ b/bec_widgets/examples/extreme/config_example.yaml
@@ -7,6 +7,35 @@ plot_settings:
 
 plot_data:
   - plot_name: "BPM plot"
+    x:
+      label: 'Motor X'
+      signals:
+        - name: "samx"
+#          entry: "samx"
+    y:
+      label: 'BPM'
+      signals:
+        - name: "gauss_bpm"
+          entry: "gauss_bpm"
+        - name: "gauss_adc1"
+          entry: "gauss_adc1"
+        - name: "gauss_adc2"
+          entry: "gauss_adc2"
+
+  - plot_name: "BPM plot 2"
+    x:
+      label: 'Motor X'
+      signals:
+        - name: "samx"
+          entry: "samx"
+    y:
+      label: 'BPM'
+      signals:
+        - name: "gauss_bpm"
+          entry: "gauss_bpm"
+        - name: "gauss_adc1"
+          entry: "gauss_adc1"
+  - plot_name: "BPM plot 3"
     x:
       label: 'Motor X'
       signals:
@@ -19,7 +48,6 @@ plot_data:
           entry: "gauss_bpm"
         - name: "gauss_adc1"
           entry: "gauss_adc1"
-
   - plot_name: "ADC plot"
     x:
       label: 'Motor Y'
diff --git a/bec_widgets/examples/extreme/extreme.py b/bec_widgets/examples/extreme/extreme.py
index 97d4f7df..a4aa53c9 100644
--- a/bec_widgets/examples/extreme/extreme.py
+++ b/bec_widgets/examples/extreme/extreme.py
@@ -74,6 +74,8 @@ class PlotApp(QWidget):
         self.grid_coordinates = None
         self.scanID = None
 
+        self.user_colors = {}  # key: (plot_name, y_name, y_entry), value: color
+
         # Initialize the UI
         self.init_ui(self.plot_settings["num_columns"])
         self.spinBox_N_columns.setValue(
@@ -200,8 +202,11 @@ class PlotApp(QWidget):
                     y_entries = [y_entries]
 
                 for y_entry in y_entries:
-                    pen_curve = mkPen(color=color, width=2, style=QtCore.Qt.DashLine)
-                    brush_curve = mkBrush(color=color)
+                    user_color = self.user_colors.get((plot_name, y_name, y_entry), None)
+                    color_to_use = user_color if user_color else color
+
+                    pen_curve = mkPen(color=color_to_use, width=2, style=QtCore.Qt.DashLine)
+                    brush_curve = mkBrush(color=color_to_use)
 
                     curve_data = pg.PlotDataItem(
                         symbolSize=5,
@@ -217,9 +222,11 @@ class PlotApp(QWidget):
 
                     # Create a ColorButton and set its color
                     color_btn = ColorButton()
-                    color_btn.setColor(color)
+                    color_btn.setColor(color_to_use)
                     color_btn.sigColorChanged.connect(
-                        lambda btn=color_btn, curve=curve_data: self.change_curve_color(btn, curve)
+                        lambda btn=color_btn, plot=plot_name, yname=y_name, yentry=y_entry, curve=curve_data: self.change_curve_color(
+                            btn, plot, yname, yentry, curve
+                        )
                     )
 
                     # Add the ColorButton as a QWidget to the table
@@ -238,13 +245,31 @@ class PlotApp(QWidget):
         self.tableWidget_crosshair.setVerticalHeaderLabels(row_labels)
         self.hook_crosshair()
 
-    def change_curve_color(self, btn, curve):
-        """Change the color of a curve."""
+    # def change_curve_color(self, btn, curve):
+    #     """Change the color of a curve."""
+    #     color = btn.color()
+    #     pen_curve = mkPen(color=color, width=2, style=QtCore.Qt.DashLine)
+    #     brush_curve = mkBrush(color=color)
+    #     curve.setPen(pen_curve)
+    #     curve.setSymbolBrush(brush_curve)
+
+    def change_curve_color(self, btn, plot_name, y_name, y_entry, curve):
+        """
+        Change the color of a curve and update the corresponding ColorButton.
+
+        Args:
+            btn (ColorButton): The ColorButton that was clicked.
+            plot_name (str): The name of the plot where the curve belongs.
+            y_name (str): The name of the y signal.
+            y_entry (str): The entry of the y signal.
+            curve (PlotDataItem): The curve to be changed.
+        """
         color = btn.color()
         pen_curve = mkPen(color=color, width=2, style=QtCore.Qt.DashLine)
         brush_curve = mkBrush(color=color)
         curve.setPen(pen_curve)
         curve.setSymbolBrush(brush_curve)
+        self.user_colors[(plot_name, y_name, y_entry)] = color
 
     def hook_crosshair(self):
         """Attach crosshairs to each plot and connect them to the update_table method."""
@@ -322,6 +347,7 @@ class PlotApp(QWidget):
             msg (dict): Message received with scan data.
             metadata (dict): Metadata of the scan.
         """
+
         current_scanID = msg.get("scanID", None)
         if current_scanID is None:
             return
@@ -452,7 +478,7 @@ if __name__ == "__main__":
 
     parser = argparse.ArgumentParser(description="Plotting App")
     parser.add_argument(
-        "--config", "-c", help="Path to the .yaml configuration file", default="config.yaml"
+        "--config", "-c", help="Path to the .yaml configuration file", default="config_example.yaml"
     )
     args = parser.parse_args()