import jwt from PySide6.QtCore import Qt, Slot, Signal from PySide6.QtGui import QAction from PySide6.QtWidgets import ( QMainWindow, QWidget, QSplitter, QHBoxLayout, QVBoxLayout, QMessageBox, QApplication, QDockWidget, QTabWidget) from aaredaqlib.coordinate import Coordinate, SmargonCoordinate from aaredaqlib.diffraction_geometry import DiffractionGeometry from aaredaqlib.logger_config import setup_logger from aaredaqlib.models import SampleShortInfoList, TokenData, DAQStatusModel, BeamlineStateEnum from aaredaqlib.sample_geometry import SampleGeometryModel from aaregui.panels.LogPanel import LogDock from aaregui.panels.beamline_controls import BeamlineControls from aaregui.panels.data_collection_settings import DataCollectionSettings from aaregui.panels.manual_sample_panel import ManualSamplePanel from aaregui.panels.reference_tools_panel import ReferenceToolsPanel from aaregui.panels.sample_queue_panel import SampleQueuePanel from aaregui.panels.tell_sample_panel import TellSamplePanel from aaregui.panels.face_detection_panel import FaceDetectionPanel from aaregui.scan_logic.raster_grid_manager import RasterGridManager from aaregui.scan_logic.rotation_scan_manager import RotationScanManager from aaregui.scan_logic.sample_mount_logic import SampleMountLogic from aaregui.threads.axis_video_thread import VideoThread from aaregui.threads.camera_thread import SampleCameraThread, PredictionSubscriber from aaregui.threads.daq_worker import DAQWorker from aaregui.threads.jfjoch_viewer import JFJochDBusClient from aaregui.widgets.camera_image import SampleCameraImageLabel from aaregui.widgets.no_wheel_scroll_area import NoWheelScrollArea from aaregui.widgets.status_bar import StatusBar from aaregui.widgets.video_image import VideoGraphicsView from aaregui.panels.fluorescence_panel import FluorescencePanel logger = setup_logger("aareGUI") class MainWindow(QMainWindow): sample_geometry = Signal(SampleGeometryModel) def __init__(self, base_url: str | None, token: str, default_image: str | None, zmq_addr: str | None, pred_zmq_addr: str | None): super().__init__() self.__base_url = base_url self.__token = token self.__mounting = False self.viewer = JFJochDBusClient() # Decode the JWT without signature verification self.__decoded_token = TokenData(**jwt.decode(token, options={"verify_signature": False})) logger.debug(self.__decoded_token) self.setStyleSheet("background-color: rgb(216, 228, 253);") top_widget = QWidget(parent=self) top_widget_layout = QHBoxLayout(top_widget) top_widget.setLayout(top_widget_layout) diffraction = DiffractionGeometry( energy_keV=12.4, dtz_mm=100, detector_size_pxl=(1553,1630), pixel_size_mm=0.150, #PILATUS 4 beam_center_pxl=(750, 750), detector_description="PILATUS 4", detector_serial_number="1", poni_rot1_rad=-0.001396263, poni_rot2_rad=-0.003839724 ) geom = SampleGeometryModel(beam_location_pxl=Coordinate(x=1000,y=1000), pixel_in_mm=0.001, aerotech=Coordinate(), smargon=SmargonCoordinate(sh_mm=Coordinate(), phi_deg=0, chi_deg=0), omega_deg=0, beam_size_mm=Coordinate(x=0.01, y=0.01), aerotech_meas=Coordinate() ) self.raster = RasterGridManager(geom=geom) self.rotation = RotationScanManager() collection_controls_scroll = NoWheelScrollArea(top_widget) self.data_collection = DataCollectionSettings(s=geom, parent=top_widget, raster_mgr=self.raster, diffraction=diffraction) top_widget_layout.addWidget(collection_controls_scroll) collection_controls_scroll.setWidget(self.data_collection) collection_controls_scroll.setHorizontalScrollBarPolicy( Qt.ScrollBarPolicy.ScrollBarAlwaysOff ) collection_controls_scroll.setFixedWidth(self.data_collection.set_width + 10) self.video_tab = QTabWidget(parent=top_widget) self.sample_camera = SampleCameraImageLabel(geom=geom, raster=self.raster, parent=top_widget, default_image=default_image) self.beamline_view_container = QWidget(parent=top_widget) self.beamline_view_layout = QVBoxLayout(self.beamline_view_container) self.beamline_view_layout.setContentsMargins(0,0,0,0) self.beamline_view_layout.setSpacing(6) self.beamline_view_1_combined = VideoGraphicsView() self.beamline_view_2_combined = VideoGraphicsView() self.beamline_view_layout.addWidget(self.beamline_view_1_combined) self.beamline_view_layout.addWidget(self.beamline_view_2_combined) self.beamline_view = VideoGraphicsView() self.beamline_camera_thread = VideoThread(ip="x06da-axis-1.psi.ch") self.beamline_camera_thread.frame_ready.connect(self.beamline_view.update_frame) self.beamline_camera_thread.frame_ready.connect(self.beamline_view_2_combined.update_frame) self.beamline_camera_thread.start() self.beamline_view_layout.addWidget(self.beamline_view) self.gonio_view = VideoGraphicsView() #TODO add option to change cameras for gonio_camera_thread self.gonio_camera_thread = VideoThread(ip="axis-accc8ed2972e.psi.ch", camera=3) self.gonio_camera_thread.frame_ready.connect(self.gonio_view.update_frame) self.gonio_camera_thread.frame_ready.connect(self.beamline_view_1_combined.update_frame) self.gonio_camera_thread.start() self.beamline_view_layout.addWidget(self.gonio_view) self.video_tab.addTab(self.sample_camera, "Sample camera") self.video_tab.addTab(self.gonio_view, "Gonio camera") self.video_tab.addTab(self.beamline_view, "Beamline view") self.video_tab.addTab(self.beamline_view_container, "Beamline combined view") top_widget_layout.addWidget(self.video_tab) beamline_controls_scroll = NoWheelScrollArea(top_widget) self.beamline = BeamlineControls(beamline_controls_scroll, staff=self.__decoded_token.staff) top_widget_layout.addWidget(beamline_controls_scroll) beamline_controls_scroll.setWidget(self.beamline) beamline_controls_scroll.setHorizontalScrollBarPolicy( Qt.ScrollBarPolicy.ScrollBarAlwaysOff ) beamline_controls_scroll.setFixedWidth(self.beamline.set_width + 10) self.tell_samples = TellSamplePanel(samples=SampleShortInfoList(s = [])) self.ref_tools_panel = ReferenceToolsPanel(samples=SampleShortInfoList(s = [])) self.job_list_panel = SampleQueuePanel() self.tell_samples_dock = QDockWidget("Sample List", self) self.tell_samples_dock.setObjectName("tell_samples_dock") self.tell_samples_dock.setWidget(self.tell_samples) self.tell_samples_dock.setAllowedAreas(Qt.DockWidgetArea.BottomDockWidgetArea) self.addDockWidget(Qt.DockWidgetArea.BottomDockWidgetArea, self.tell_samples_dock) self.ref_tools_dock = QDockWidget("Reference Tools", self) self.ref_tools_dock.setWidget(self.ref_tools_panel) self.ref_tools_dock.setAllowedAreas(Qt.DockWidgetArea.BottomDockWidgetArea) self.addDockWidget(Qt.DockWidgetArea.BottomDockWidgetArea, self.ref_tools_dock) self.tabifyDockWidget(self.ref_tools_dock, self.tell_samples_dock) self.sample_logic = SampleMountLogic() self.job_list_dock = QDockWidget("Automation list", self) self.job_list_dock.setWidget(self.job_list_panel) self.job_list_dock.setAllowedAreas(Qt.DockWidgetArea.BottomDockWidgetArea) self.addDockWidget(Qt.DockWidgetArea.BottomDockWidgetArea, self.job_list_dock) self.manual_sample_panel = ManualSamplePanel() self.manual_sample_dock = QDockWidget("Manual sample", self) self.manual_sample_dock.setWidget(self.manual_sample_panel) self.manual_sample_dock.setAllowedAreas(Qt.DockWidgetArea.BottomDockWidgetArea) self.addDockWidget(Qt.DockWidgetArea.BottomDockWidgetArea, self.manual_sample_dock) self.face_panel = FaceDetectionPanel() self.face_panel_dock = QDockWidget("Face detection", self) self.face_panel_dock.setWidget(self.face_panel) self.face_panel_dock.setAllowedAreas(Qt.DockWidgetArea.RightDockWidgetArea | Qt.DockWidgetArea.LeftDockWidgetArea) self.addDockWidget(Qt.DockWidgetArea.RightDockWidgetArea, self.face_panel_dock) self.face_panel_dock.hide() self.fluor_panel = FluorescencePanel() self.fluor_panel_dock = QDockWidget("Fluorescence", self) self.fluor_panel_dock.setWidget(self.fluor_panel) self.fluor_panel_dock.setAllowedAreas(Qt.DockWidgetArea.TopDockWidgetArea | Qt.DockWidgetArea.BottomDockWidgetArea | Qt.DockWidgetArea.RightDockWidgetArea) self.addDockWidget(Qt.DockWidgetArea.BottomDockWidgetArea, self.fluor_panel_dock) self.fluor_panel_dock.hide() # Create and add the dock to your main window self.log_dock = LogDock("Console Log", self) self.addDockWidget(Qt.BottomDockWidgetArea, self.log_dock) self.log_dock.attach_logger("") self.log_dock.attach_logger("aareDAQ") self.log_dock.attach_logger("aareGUI") self.log_dock.hide() self.setCentralWidget(top_widget) self.setWindowTitle("AareGUI") self.create_menu_bar() self.status_bar = StatusBar(self.__decoded_token, parent=self) self.setStatusBar(self.status_bar) self.daq = DAQWorker(base_url=self.__base_url, token=self.__token) self.daq.spreadsheet.connect(self.tell_samples.new_sample_list) if self.__decoded_token.staff: self.daq.reference_tools.connect(self.ref_tools_panel.new_list) self.beamline.samcam.changed.connect(self.daq.samcam_settings) self.beamline.loopctr.background.clicked.connect(self.daq.alc_background) self.beamline.loopctr.find_tip.clicked.connect(self.daq.center_loop) self.beamline.loopctr.bounding_box.clicked.connect(self.daq.ml_bounding_box) self.daq.raster_generated_by_ml.connect(self.raster.update_active_grid_request) self.sample_camera.smargon.connect(self.daq.move_smargon) self.beamline.smargon_panel.smargon.connect(self.daq.move_smargon) self.sample_camera.samcam_updated.connect(self.daq.samcam_settings) self.raster.omega.connect(self.daq.set_omega) self.raster.smargon.connect(self.daq.move_smargon) self.beamline.omega_panel.set_omega.connect(self.daq.set_omega) self.sample_camera.set_omega.connect(self.daq.set_omega) self.beamline.zoom_panel.zoom.connect(self.daq.zoom) self.beamline.illumination_panel.light.connect(self.daq.light) if self.__decoded_token.staff: self.beamline.abr_tweak.abr_tweak.connect(self.daq.abr_tweak) self.beamline.abr_tweak.abr_save.connect(self.daq.abr_save) self.beamline.abr_tweak.abr_goto_meas.connect(self.daq.abr_goto_meas) self.beamline.beam_mark.beam_mark_clear.connect(self.daq.beam_mark_clear) self.sample_camera.update_beam_mark.connect(self.daq.beam_mark_add) self.beamline.beam_center.beam_center.connect(self.daq.beam_center) self.beamline.beam_size.beam_size.connect(self.daq.beam_size_mm) self.sample_camera.load_image.connect(self.raster.load_image) self.sample_camera.switch_raster_grid.connect(self.data_collection.switch_to_raster) self.beamline.samcam.show_detections_changed.connect(self.sample_camera.set_show_detections) if zmq_addr is not None: self.camera_thread = SampleCameraThread(zmq_url=zmq_addr) self.camera_thread.camera_image.connect(self.sample_camera.update_pixmap) self.camera_thread.start() else: self.camera_thread = None # Prediction subscriber thread if pred_zmq_addr is not None: logger.debug(f"Starting prediction subscriber thread {pred_zmq_addr}") self.prediction_thread = PredictionSubscriber(pred_zmq_url=pred_zmq_addr) self.prediction_thread.prediction.connect(self.sample_camera.update_detections) self.prediction_thread.start() else: self.prediction_thread = None QApplication.instance().aboutToQuit.connect(self.cleanup) # # self.data_collection.helical.helical_scan.connect(self.worker.helical_scan) # self.data_collection.helical.update_bookmarks.connect( # self.camera_image.update_bookmarks # ) # self.sample_camera.autofocus.connect(self.daq.autofocus) self.sample_camera.evaluate_grid.connect(self.raster.run_grid_scan) self.data_collection.raster.evaluate_grid.connect(self.raster.run_grid_scan) self.data_collection.raster.evaluate_grid_auto.connect(self.raster.run_grid_scan_auto) self.sample_camera.clear_evaluated_grids.connect(self.raster.clear_completed_grids) self.sample_camera.clear_grid.connect(self.raster.clear_active_grid) self.daq.run_number_incremented.connect(self.data_collection.file_path_panel.increment_run_number) self.job_list_panel.auto_scan.connect(self.daq.automated_scan) self.job_list_panel.unmount.connect(self.daq.unmount) self.tell_samples.mount.connect(self.daq.mount) self.tell_samples.unmount.connect(self.daq.unmount) self.ref_tools_panel.mount.connect(self.daq.mount) self.ref_tools_panel.unmount.connect(self.daq.unmount) self.data_collection.raster.grid_size_updated.connect(self.raster.update_grid_size) self.data_collection.raster.exp_time_updated.connect(self.raster.update_exposure_time) self.data_collection.raster.transmission_updated.connect(self.raster.update_transmission) self.data_collection.raster.dtz_updated.connect(self.raster.update_dtz) self.data_collection.raster.grid_metric_updated.connect(self.raster.metric) self.data_collection.raster.raster_alpha_changed.connect(self.sample_camera.raster_alpha) self.data_collection.cancel.connect(self.daq.cancel) self.raster.grid_scan.connect(self.daq.raster_scan) self.raster.grid_scan_auto.connect(self.daq.raster_scan_auto) self.data_collection.screening.rotation_scan.connect(self.daq.standard_scan) self.data_collection.simple.rotation_scan.connect(self.daq.standard_scan) self.data_collection.simple.parameters_changed.connect(self.daq.smart_params) self.raster.grid_scan_size_changed.connect(self.data_collection.raster.grid_scan_size_change) self.status_bar.set_pgroup.connect(self.daq.set_pgroup) self.status_bar.end_session.connect(self.daq.end_session) self.status_bar.force_session.connect(self.daq.force_session) self.status_bar.dewar_exchange.connect(self.daq.dewar_exchange) self.status_bar.sample_exchange.connect(self.daq.sample_exchange) self.status_bar.sample_alignment.connect(self.daq.sample_alignment) self.status_bar.beam_location.connect(self.daq.beam_location) self.status_bar.close_shutter.connect(self.daq.close_shutter) self.status_bar.open_shutter.connect(self.daq.open_shutter) self.rotation.file_ready.connect(self.viewer.load_image) self.raster.image_selected.connect(self.viewer.load_image) self.raster.viewer_track_online.connect(self.viewer.load_online) self.data_collection.screening.viewer_track_online.connect(self.viewer.load_online) self.job_list_panel.viewer_track_online.connect(self.viewer.load_online) self.sample_logic.sample_changed.connect(self.data_collection.file_path_panel.update_sample) self.daq.update.connect(self.beamline.omega_panel.update_daq_status) self.daq.update.connect(self.beamline.smargon_panel.update_daq_status) self.daq.update.connect(self.job_list_panel.update_daq_status) self.daq.update.connect(self.data_collection.file_path_panel.update_daq_status) self.daq.update.connect(self.data_collection.update_daq_status) self.daq.update.connect(self.beamline.illumination_panel.update_daq_status) self.daq.update.connect(self.raster.update_daq_status) self.daq.update.connect(self.sample_camera.update_daq_status) self.daq.update.connect(self.tell_samples.update_daq_status) self.daq.update.connect(self.ref_tools_panel.update_daq_status) if self.__decoded_token.staff: self.daq.update.connect(self.beamline.beam_size.update_daq_status) self.daq.update.connect(self.beamline.beam_center.update_daq_status) self.daq.update.connect(self.beamline.abr_tweak.update_daq_status) self.daq.update.connect(self.beamline.beam_mark.update_daq_status) self.daq.update.connect(self.status_bar.update_daq_status) self.daq.update.connect(self.beamline.samcam.update_daq_status) self.daq.update.connect(self.beamline.zoom_panel.update_daq_status) self.daq.update.connect(self.sample_logic.update_daq_status) self.daq.update.connect(self.manual_sample_panel.update_daq_status) self.daq.update.connect(self.fluor_panel.update_daq_status) self.daq.update.connect(self.update_daq_status) self.daq.sample_missing.connect(self.show_sample_missing_dialog) self.daq.standard_scan_completed.connect(self.rotation.scan_completed) self.daq.automated_scan_done.connect(self.job_list_panel.automated_scan_done) self.daq.raster_scan_completed.connect(self.raster.grid_scan_completed) self.manual_sample_panel.sample_manual.connect(self.daq.sample_manual) self.face_panel.face_detection.connect(self.daq.face_detection) self.daq.face_detection_result.connect(self.face_panel.update_plot) self.status_bar.get_all_pgroups.connect(self.daq.get_all_pgroups) self.daq.staff_pgroups_loaded.connect(self.status_bar.staff_pgroups_loaded) self.data_collection.fluo.fluo_scan.connect(self.daq.fluorimeter_spectrum) self.daq.fluorimeter_spectrum_update.connect(self.fluor_panel.update_plot) self.daq.fluorimeter_spectrum_update.connect(lambda: self.fluor_panel_dock.setVisible(True)) def create_menu_bar(self): """Create a menu bar with File->Quit and Help->About.""" # Main menu bar menu_bar = self.menuBar() # File menu file_menu = menu_bar.addMenu("File") quit_action = QAction("&Quit", self) # Action for 'Quit' quit_action.triggered.connect(self.close) # Connect to window's `close` method file_menu.addAction(quit_action) view_menu = menu_bar.addMenu("View") show_samples_action = QAction("Show Sample List", self) show_samples_action.setCheckable(True) show_samples_action.setChecked(True) show_samples_action.triggered.connect(lambda: self.tell_samples_dock.setVisible(show_samples_action.isChecked())) # When dock visibility changes, update the action's checked state self.tell_samples_dock.visibilityChanged.connect(show_samples_action.setChecked) view_menu.addAction(show_samples_action) show_job_list_action = QAction("Show job List", self) show_job_list_action.setCheckable(True) show_job_list_action.setChecked(True) show_job_list_action.triggered.connect(self.job_list_dock.setVisible) # When dock visibility changes, update the action's checked state show_job_list_action.triggered.connect(lambda: self.job_list_dock.setVisible(show_job_list_action.isChecked())) view_menu.addAction(show_job_list_action) show_manual_sample_action = QAction("Show manual sample", self) show_manual_sample_action.setCheckable(True) show_manual_sample_action.setChecked(True) show_manual_sample_action.triggered.connect(self.manual_sample_dock.setVisible) # When dock visibility changes, update the action's checked state show_manual_sample_action.triggered.connect(lambda: self.manual_sample_dock.setVisible(show_manual_sample_action.isChecked())) view_menu.addAction(show_manual_sample_action) show_face_panel_action = QAction("Show face detection", self) show_face_panel_action.setCheckable(True) show_face_panel_action.setChecked(False) #show_face_panel_action.triggered.connect(self.face_panel_dock.setVisible) show_face_panel_action.triggered.connect(lambda checked: self.face_panel_dock.setVisible(checked)) self.face_panel_dock.visibilityChanged.connect(show_face_panel_action.setChecked) view_menu.addAction(show_face_panel_action) show_fluor_panel_action = QAction("Show fluorescence", self) show_fluor_panel_action.setCheckable(True) show_fluor_panel_action.setChecked(False) show_fluor_panel_action.triggered.connect(lambda checked: self.fluor_panel_dock.setVisible(checked)) self.fluor_panel_dock.visibilityChanged.connect(show_fluor_panel_action.setChecked) view_menu.addAction(show_fluor_panel_action) show_log_action = QAction("Show Log", self) show_log_action.setCheckable(True) show_log_action.setChecked(False) show_log_action.triggered.connect(lambda checked: self.log_dock.setVisible(checked)) self.log_dock.visibilityChanged.connect(show_log_action.setChecked) view_menu.addAction(show_log_action) help_menu = menu_bar.addMenu("Help") about_action = QAction("About", self) # Action for 'About' about_action.triggered.connect(self.show_about_dialog) help_menu.addAction(about_action) def show_about_dialog(self): QMessageBox.about( self, "About", "Aare Macromolecular Crystallography GUI\nVersion: 1.0\nCopyright: Paul Scherrer Institute 2024-2025", ) @Slot(str) def show_sample_missing_dialog(self, msg: str): if self.job_list_panel.is_running(): logger.warning(f"Automation: Sample missing: {msg}") return QMessageBox.warning(self, "No Sample", f"{msg}") #TODO tidy up mount and sampel view fucntions @Slot() def mount_view(self): self.video_tab.setCurrentIndex(3) @Slot() def sample_view(self): self.video_tab.setCurrentIndex(0) @Slot(DAQStatusModel) def update_daq_status(self, s: DAQStatusModel): self.beamline_camera_thread.set_busy(s.busy) self.gonio_camera_thread.set_busy(s.busy) if not self.__mounting and s.state == BeamlineStateEnum.RobotSampleExchange: self.__mounting = True self.video_tab.setCurrentIndex(3) elif self.__mounting and s.state != BeamlineStateEnum.RobotSampleExchange: self.__mounting = False self.video_tab.setCurrentIndex(0) @Slot(str) def display_error(self, msg: str): self.status_bar.setStyleSheet("color: red;") self.status_bar.showMessage(f' {msg} ', 10000) def cleanup(self): if self.camera_thread is not None: self.camera_thread.stop() if self.prediction_thread is not None: self.prediction_thread.stop() if self.beamline_camera_thread is not None: self.beamline_camera_thread.stop() if self.gonio_camera_thread is not None: self.gonio_camera_thread.stop()