From c4b85381a41e4742567680864668ee83d498b1d1 Mon Sep 17 00:00:00 2001
From: wakonig_k <klaus.wakonig@psi.ch>
Date: Tue, 10 Sep 2024 09:53:58 +0200
Subject: [PATCH] feat(scan_progressbar): added progressbar with hooks to scan
 progress and device progress

---
 bec_widgets/cli/client.py                     |  10 +
 .../progress/scan_progressbar/__init__.py     |   0
 .../register_scan_progress_bar.py             |  17 +
 .../scan_progress_bar.pyproject               |   1 +
 .../scan_progress_bar_plugin.py               |  54 ++++
 .../scan_progressbar/scan_progressbar.py      | 298 ++++++++++++++++++
 .../scan_progressbar/scan_progressbar.ui      | 141 +++++++++
 7 files changed, 521 insertions(+)
 create mode 100644 bec_widgets/widgets/progress/scan_progressbar/__init__.py
 create mode 100644 bec_widgets/widgets/progress/scan_progressbar/register_scan_progress_bar.py
 create mode 100644 bec_widgets/widgets/progress/scan_progressbar/scan_progress_bar.pyproject
 create mode 100644 bec_widgets/widgets/progress/scan_progressbar/scan_progress_bar_plugin.py
 create mode 100644 bec_widgets/widgets/progress/scan_progressbar/scan_progressbar.py
 create mode 100644 bec_widgets/widgets/progress/scan_progressbar/scan_progressbar.ui

diff --git a/bec_widgets/cli/client.py b/bec_widgets/cli/client.py
index 1c282c86..06f8d87c 100644
--- a/bec_widgets/cli/client.py
+++ b/bec_widgets/cli/client.py
@@ -3245,6 +3245,16 @@ class ScanControl(RPCBase):
         """
 
 
+class ScanProgressBar(RPCBase):
+    """Widget to display a progress bar that is hooked up to the scan progress of a scan."""
+
+    @rpc_call
+    def remove(self):
+        """
+        Cleanup the BECConnector
+        """
+
+
 class ScatterCurve(RPCBase):
     """Scatter curve item for the scatter waveform widget."""
 
diff --git a/bec_widgets/widgets/progress/scan_progressbar/__init__.py b/bec_widgets/widgets/progress/scan_progressbar/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/bec_widgets/widgets/progress/scan_progressbar/register_scan_progress_bar.py b/bec_widgets/widgets/progress/scan_progressbar/register_scan_progress_bar.py
new file mode 100644
index 00000000..4d61e81b
--- /dev/null
+++ b/bec_widgets/widgets/progress/scan_progressbar/register_scan_progress_bar.py
@@ -0,0 +1,17 @@
+def main():  # pragma: no cover
+    from qtpy import PYSIDE6
+
+    if not PYSIDE6:
+        print("PYSIDE6 is not available in the environment. Cannot patch designer.")
+        return
+    from PySide6.QtDesigner import QPyDesignerCustomWidgetCollection
+
+    from bec_widgets.widgets.progress.scan_progressbar.scan_progress_bar_plugin import (
+        ScanProgressBarPlugin,
+    )
+
+    QPyDesignerCustomWidgetCollection.addCustomWidget(ScanProgressBarPlugin())
+
+
+if __name__ == "__main__":  # pragma: no cover
+    main()
diff --git a/bec_widgets/widgets/progress/scan_progressbar/scan_progress_bar.pyproject b/bec_widgets/widgets/progress/scan_progressbar/scan_progress_bar.pyproject
new file mode 100644
index 00000000..80d55e87
--- /dev/null
+++ b/bec_widgets/widgets/progress/scan_progressbar/scan_progress_bar.pyproject
@@ -0,0 +1 @@
+{'files': ['scan_progressbar.py']}
\ No newline at end of file
diff --git a/bec_widgets/widgets/progress/scan_progressbar/scan_progress_bar_plugin.py b/bec_widgets/widgets/progress/scan_progressbar/scan_progress_bar_plugin.py
new file mode 100644
index 00000000..5c8bf71a
--- /dev/null
+++ b/bec_widgets/widgets/progress/scan_progressbar/scan_progress_bar_plugin.py
@@ -0,0 +1,54 @@
+# Copyright (C) 2022 The Qt Company Ltd.
+# SPDX-License-Identifier: LicenseRef-Qt-Commercial OR BSD-3-Clause
+
+from qtpy.QtDesigner import QDesignerCustomWidgetInterface
+
+from bec_widgets.utils.bec_designer import designer_material_icon
+from bec_widgets.widgets.progress.scan_progressbar.scan_progressbar import ScanProgressBar
+
+DOM_XML = """
+<ui language='c++'>
+    <widget class='ScanProgressBar' name='scan_progress_bar'>
+    </widget>
+</ui>
+"""
+
+
+class ScanProgressBarPlugin(QDesignerCustomWidgetInterface):  # pragma: no cover
+    def __init__(self):
+        super().__init__()
+        self._form_editor = None
+
+    def createWidget(self, parent):
+        t = ScanProgressBar(parent)
+        return t
+
+    def domXml(self):
+        return DOM_XML
+
+    def group(self):
+        return "BEC Utils"
+
+    def icon(self):
+        return designer_material_icon(ScanProgressBar.ICON_NAME)
+
+    def includeFile(self):
+        return "scan_progress_bar"
+
+    def initialize(self, form_editor):
+        self._form_editor = form_editor
+
+    def isContainer(self):
+        return False
+
+    def isInitialized(self):
+        return self._form_editor is not None
+
+    def name(self):
+        return "ScanProgressBar"
+
+    def toolTip(self):
+        return "A progress bar that is hooked up to the scan progress of a scan."
+
+    def whatsThis(self):
+        return self.toolTip()
diff --git a/bec_widgets/widgets/progress/scan_progressbar/scan_progressbar.py b/bec_widgets/widgets/progress/scan_progressbar/scan_progressbar.py
new file mode 100644
index 00000000..1ba97cb4
--- /dev/null
+++ b/bec_widgets/widgets/progress/scan_progressbar/scan_progressbar.py
@@ -0,0 +1,298 @@
+from __future__ import annotations
+
+import enum
+import os
+import time
+
+import numpy as np
+from bec_lib.endpoints import MessageEndpoints
+from bec_lib.logger import bec_logger
+from qtpy.QtCore import QObject, QTimer
+from qtpy.QtWidgets import QVBoxLayout, QWidget
+
+from bec_widgets.utils.bec_widget import BECWidget
+from bec_widgets.utils.error_popups import SafeProperty, SafeSlot
+from bec_widgets.utils.ui_loader import UILoader
+
+logger = bec_logger.logger
+
+
+class ProgressSource(enum.Enum):
+    """
+    Enum to define the source of the progress.
+    """
+
+    SCAN_PROGRESS = "scan_progress"
+    DEVICE_PROGRESS = "device_progress"
+
+
+class ProgressTask(QObject):
+    """
+    Class to store progress information.
+    Inspired by https://github.com/Textualize/rich/blob/master/rich/progress.py
+    """
+
+    def __init__(self, parent: QWidget, value: float = 0, max_value: float = 0, done: bool = False):
+        super().__init__(parent=parent)
+        self.start_time = time.time()
+        self.done = done
+        self.value = value
+        self.max_value = max_value
+        self._elapsed_time = 0
+
+        self.timer = QTimer(self)
+        self.timer.timeout.connect(self.update_elapsed_time)
+        self.timer.start(100)  # update the elapsed time every 100 ms
+
+    def update(self, value: float, max_value: float, done: bool = False):
+        """
+        Update the progress.
+        """
+        self.max_value = max_value
+        self.done = done
+        self.value = value
+        if done:
+            self.timer.stop()
+
+    def update_elapsed_time(self):
+        """
+        Update the time estimates. This is called every 100 ms by a QTimer.
+        """
+        self._elapsed_time += 0.1
+
+    @property
+    def percentage(self) -> float:
+        """float: Get progress of task as a percentage. If a None total was set, returns 0"""
+        if not self.max_value:
+            return 0.0
+        completed = (self.value / self.max_value) * 100.0
+        completed = min(100.0, max(0.0, completed))
+        return completed
+
+    @property
+    def speed(self) -> float:
+        """Get the estimated speed in steps per second."""
+        if self._elapsed_time == 0:
+            return 0.0
+
+        return self.value / self._elapsed_time
+
+    @property
+    def frequency(self) -> float:
+        """Get the estimated frequency in steps per second."""
+        if self.speed == 0:
+            return 0.0
+        return 1 / self.speed
+
+    @property
+    def time_elapsed(self) -> str:
+        # format the elapsed time to a string in the format HH:MM:SS
+        return self._format_time(int(self._elapsed_time))
+
+    @property
+    def remaining(self) -> float:
+        """Get the estimated remaining steps."""
+        if self.done:
+            return 0.0
+        remaining = self.max_value - self.value
+        return remaining
+
+    @property
+    def time_remaining(self) -> str:
+        """
+        Get the estimated remaining time in the format HH:MM:SS.
+        """
+        if self.done or not self.speed or not self.remaining:
+            return self._format_time(0)
+        estimate = int(np.round(self.remaining / self.speed))
+
+        return self._format_time(estimate)
+
+    def _format_time(self, seconds: float) -> str:
+        """
+        Format the time in seconds to a string in the format HH:MM:SS.
+        """
+        return f"{seconds // 3600:02}:{(seconds // 60) % 60:02}:{seconds % 60:02}"
+
+
+class ScanProgressBar(BECWidget, QWidget):
+    """
+    Widget to display a progress bar that is hooked up to the scan progress of a scan.
+    If you want to manually set the progress, it is recommended to use the BECProgressbar or QProgressbar directly.
+    """
+
+    ICON_NAME = "timelapse"
+
+    def __init__(self, parent=None, client=None, config=None, gui_id=None):
+        super().__init__(parent=parent, client=client, config=config, gui_id=gui_id)
+
+        self.get_bec_shortcuts()
+        ui_file = os.path.join(os.path.dirname(__file__), "scan_progressbar.ui")
+        self.ui = UILoader(self).loader(ui_file)
+        self.layout = QVBoxLayout(self)
+        self.layout.setSpacing(0)
+        self.layout.setContentsMargins(0, 0, 0, 0)
+        self.layout.addWidget(self.ui)
+        self.setLayout(self.layout)
+        self.progressbar = self.ui.progressbar
+
+        self.connect_to_queue()
+        self._progress_source = None
+        self.task = None
+        self.scan_number = None
+
+    def connect_to_queue(self):
+        """
+        Connect to the queue status signal.
+        """
+        self.bec_dispatcher.connect_slot(self.on_queue_update, MessageEndpoints.scan_queue_status())
+
+    def set_progress_source(self, source: ProgressSource, device=None):
+        """
+        Set the source of the progress.
+        """
+        if self._progress_source == source:
+            self.update_source_label(source, device=device)
+            return
+        if self._progress_source is not None:
+            self.bec_dispatcher.disconnect_slot(
+                self.on_progress_update,
+                (
+                    MessageEndpoints.scan_progress()
+                    if self._progress_source == ProgressSource.SCAN_PROGRESS
+                    else MessageEndpoints.device_progress(device=device)
+                ),
+            )
+        self._progress_source = source
+        self.bec_dispatcher.connect_slot(
+            self.on_progress_update,
+            (
+                MessageEndpoints.scan_progress()
+                if source == ProgressSource.SCAN_PROGRESS
+                else MessageEndpoints.device_progress(device=device)
+            ),
+        )
+        self.update_source_label(source, device=device)
+
+    def update_source_label(self, source: ProgressSource, device=None):
+        scan_text = f"Scan {self.scan_number}" if self.scan_number is not None else "Scan"
+        text = scan_text if source == ProgressSource.SCAN_PROGRESS else f"Device {device}"
+        logger.info(f"Set progress source to {text}")
+        self.ui.source_label.setText(text)
+
+    def on_progress_update(self, msg_content, metadata):
+        """
+        Update the progress bar based on the progress message.
+        """
+        value = msg_content["value"]
+        max_value = msg_content.get("max_value", 100)
+        done = msg_content.get("done", False)
+
+        if self.task is None:
+            return
+        self.task.update(value, max_value, done)
+
+        self.update_labels()
+
+        self.progressbar.set_maximum(self.task.max_value)
+        if done:
+            self.progressbar.set_value(self.task.max_value)
+            self.task = None
+            return
+
+        self.progressbar.set_value(self.task.value)
+
+    @SafeProperty(bool)
+    def show_elapsed_time(self):
+        return self.ui.elapsed_time_label.isVisible()
+
+    @show_elapsed_time.setter
+    def show_elapsed_time(self, value):
+        self.ui.elapsed_time_label.setVisible(value)
+
+    @SafeProperty(bool)
+    def show_remaining_time(self):
+        return self.ui.remaining_time_label.isVisible()
+
+    @show_remaining_time.setter
+    def show_remaining_time(self, value):
+        self.ui.remaining_time_label.setVisible(value)
+
+    @SafeProperty(bool)
+    def show_source_label(self):
+        return self.ui.source_label.isVisible()
+
+    @show_source_label.setter
+    def show_source_label(self, value):
+        self.ui.source_label.setVisible(value)
+
+    def update_labels(self):
+        """
+        Update the labels based on the progress task.
+        """
+        if self.task is None:
+            return
+
+        self.ui.elapsed_time_label.setText(self.task.time_elapsed)
+        self.ui.remaining_time_label.setText(self.task.time_remaining)
+
+    @SafeSlot(dict, dict, verify_sender=True)
+    def on_queue_update(self, msg_content, metadata):
+        """
+        Update the progress bar based on the queue status.
+        """
+        if not "queue" in msg_content:
+            return
+        primary_queue_info = msg_content["queue"].get("primary", {}).get("info", [])
+        if len(primary_queue_info) == 0:
+            return
+        scan_info = primary_queue_info[0]
+        if scan_info is None:
+            return
+        if scan_info.get("status").lower() == "running" and self.task is None:
+            self.task = ProgressTask(parent=self)
+
+        active_request_block = scan_info.get("active_request_block", {})
+        if active_request_block is None:
+            return
+
+        self.scan_number = active_request_block.get("scan_number")
+        report_instructions = active_request_block.get("report_instructions", [])
+        if not report_instructions:
+            return
+
+        # for now, let's just use the first instruction
+        instruction = report_instructions[0]
+
+        if "scan_progress" in instruction:
+            self.set_progress_source(ProgressSource.SCAN_PROGRESS)
+        elif "device_progress" in instruction:
+            device = instruction["device_progress"][0]
+            self.set_progress_source(ProgressSource.DEVICE_PROGRESS, device=device)
+
+    def cleanup(self):
+        if self.task is not None:
+            self.task.timer.stop()
+            self.close()
+            self.deleteLater()
+        if self._progress_source is not None:
+            self.bec_dispatcher.disconnect_slot(
+                self.on_progress_update,
+                (
+                    MessageEndpoints.scan_progress()
+                    if self._progress_source == ProgressSource.SCAN_PROGRESS
+                    else MessageEndpoints.device_progress(device=self._progress_source.value)
+                ),
+            )
+
+
+if __name__ == "__main__":  # pragma: no cover
+    from qtpy.QtWidgets import QApplication
+
+    bec_logger.disabled_modules = ["bec_lib"]
+    app = QApplication([])
+
+    widget = ScanProgressBar()
+    widget.show()
+
+    app.exec_()
diff --git a/bec_widgets/widgets/progress/scan_progressbar/scan_progressbar.ui b/bec_widgets/widgets/progress/scan_progressbar/scan_progressbar.ui
new file mode 100644
index 00000000..a6b959e9
--- /dev/null
+++ b/bec_widgets/widgets/progress/scan_progressbar/scan_progressbar.ui
@@ -0,0 +1,141 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<ui version="4.0">
+ <class>Form</class>
+ <widget class="QWidget" name="Form">
+  <property name="geometry">
+   <rect>
+    <x>0</x>
+    <y>0</y>
+    <width>211</width>
+    <height>60</height>
+   </rect>
+  </property>
+  <property name="sizePolicy">
+   <sizepolicy hsizetype="Preferred" vsizetype="Preferred">
+    <horstretch>0</horstretch>
+    <verstretch>0</verstretch>
+   </sizepolicy>
+  </property>
+  <property name="minimumSize">
+   <size>
+    <width>0</width>
+    <height>60</height>
+   </size>
+  </property>
+  <property name="maximumSize">
+   <size>
+    <width>16777215</width>
+    <height>16777215</height>
+   </size>
+  </property>
+  <property name="windowTitle">
+   <string>Form</string>
+  </property>
+  <layout class="QVBoxLayout" name="verticalLayout">
+   <property name="spacing">
+    <number>1</number>
+   </property>
+   <property name="leftMargin">
+    <number>2</number>
+   </property>
+   <property name="topMargin">
+    <number>0</number>
+   </property>
+   <property name="rightMargin">
+    <number>2</number>
+   </property>
+   <property name="bottomMargin">
+    <number>0</number>
+   </property>
+   <item>
+    <layout class="QHBoxLayout" name="source_layout">
+     <item>
+      <widget class="QLabel" name="source_label">
+       <property name="maximumSize">
+        <size>
+         <width>16777215</width>
+         <height>20</height>
+        </size>
+       </property>
+       <property name="text">
+        <string>Scan</string>
+       </property>
+      </widget>
+     </item>
+     <item>
+      <spacer name="source_spacer">
+       <property name="orientation">
+        <enum>Qt::Orientation::Horizontal</enum>
+       </property>
+       <property name="sizeHint" stdset="0">
+        <size>
+         <width>40</width>
+         <height>10</height>
+        </size>
+       </property>
+      </spacer>
+     </item>
+    </layout>
+   </item>
+   <item>
+    <widget class="BECProgressBar" name="progressbar">
+     <property name="padding_left_right" stdset="0">
+      <double>2.000000000000000</double>
+     </property>
+    </widget>
+   </item>
+   <item>
+    <layout class="QHBoxLayout" name="timer_layout">
+     <item>
+      <widget class="QLabel" name="remaining_time_label">
+       <property name="maximumSize">
+        <size>
+         <width>16777215</width>
+         <height>20</height>
+        </size>
+       </property>
+       <property name="text">
+        <string>0:00:00</string>
+       </property>
+      </widget>
+     </item>
+     <item>
+      <spacer name="timer_spacer">
+       <property name="orientation">
+        <enum>Qt::Orientation::Horizontal</enum>
+       </property>
+       <property name="sizeHint" stdset="0">
+        <size>
+         <width>40</width>
+         <height>0</height>
+        </size>
+       </property>
+      </spacer>
+     </item>
+     <item>
+      <widget class="QLabel" name="elapsed_time_label">
+       <property name="maximumSize">
+        <size>
+         <width>16777215</width>
+         <height>20</height>
+        </size>
+       </property>
+       <property name="text">
+        <string>0:00:00</string>
+       </property>
+      </widget>
+     </item>
+    </layout>
+   </item>
+  </layout>
+ </widget>
+ <customwidgets>
+  <customwidget>
+   <class>BECProgressBar</class>
+   <extends>QWidget</extends>
+   <header>bec_progress_bar</header>
+  </customwidget>
+ </customwidgets>
+ <resources/>
+ <connections/>
+</ui>