From b212f0d081e9aa64f5fb9c6d07c6f938e6bf3337 Mon Sep 17 00:00:00 2001
From: appleb_m <martin.appleby@psi.ch>
Date: Tue, 28 Apr 2026 14:28:53 +0200
Subject: [PATCH] workflows: allow xtal_snapshot to go to any state

---
 src/aare/daq/daq.py       | 11 +++++++++--
 src/aare/daq/workflows.py | 38 +++++++++++++++++++++-----------------
 2 files changed, 30 insertions(+), 19 deletions(-)

diff --git a/src/aare/daq/daq.py b/src/aare/daq/daq.py
index ea3a94e3..06e658f9 100644
--- a/src/aare/daq/daq.py
+++ b/src/aare/daq/daq.py
@@ -2972,10 +2972,17 @@ class AareDAQ:
                         if target == BeamlineStateEnum.SampleAlignment:
                             workflows.xtal_snapshot2sa(self.__devs, self.__cfg)
                         elif target == BeamlineStateEnum.SampleExchange:
-                            workflows.xtal_snapshot2sa(self.__devs, self.__cfg)
-                            workflows.sa2se(self.__devs, self.__cfg)
+                            workflows.xtal_snapshot2se(self.__devs, self.__cfg)
+                        elif target == BeamlineStateEnum.RobotSampleExchange:
+                            workflows.xtal_snapshot2rse(self.__devs, self.__cfg)
                         elif target == BeamlineStateEnum.DataCollection:
                             workflows.xtal_snapshot2dc(self.__devs, self.__cfg)
+                        elif target == BeamlineStateEnum.XrayFluorescence:
+                            workflows.xtal_snapshot2xrf(self.__devs, self.__cfg)
+                        elif target == BeamlineStateEnum.DewarTransfer:
+                            workflows.xtal_snapshot2dh(self.__devs, self.__cfg)
+                        elif target == BeamlineStateEnum.BeamLocation:
+                            workflows.xtal_snapshot2bl(self.__devs, self.__cfg)
                         else:
                             raise TransformationInvalidException()
                 self.__cfg.state = target
diff --git a/src/aare/daq/workflows.py b/src/aare/daq/workflows.py
index bda7fcae..afb1f023 100644
--- a/src/aare/daq/workflows.py
+++ b/src/aare/daq/workflows.py
@@ -57,19 +57,6 @@ def m2se(devs: BeamlineDevices, cfg: BeamlineConfig):
     pass
 
 def sa2se(devs: BeamlineDevices, cfg: BeamlineConfig):
-    # print('move collimator and scintillator to 20')
-    # devs.collimator = 20.0
-    # devs.scintillator = 20.0
-    # print('move smargon home')
-    # devs.smargon_move_home()
-    # print('try to move aerotech')
-    # devs.aerotech_pos = ABR_POS_MOUNT
-    # print('move bl to park')
-    # devs.reflector_up = StagePositionEnum.PARK
-    # print('move bs to park')
-    # devs.beamstop_stage_up = StagePositionEnum.PARK
-    # #print('move cryo to park')
-    # #devs.__cryojet_pos = StagePositionEnum.PARK
     try:
         devs.samcam_auto(AutoEnum.AUTO)
         devs.bec_worker.planner.move_to(BeamlineState.MANUAL_SAMPLE_EXCHANGE)
@@ -77,7 +64,7 @@ def sa2se(devs: BeamlineDevices, cfg: BeamlineConfig):
         devs.aerotech_pos = cfg.abr_meas_pos
         devs.smargon_move_home()
         devs.samcam_auto(AutoEnum.ONCE)
-    except Exception as e:
+    except Exception:
         devs.aerotech_pos = cfg.abr_meas_pos
         devs.smargon_move_home()
         raise
@@ -115,6 +102,22 @@ def xtal_snapshot2sa(devs:BeamlineDevices, cfg:BeamlineConfig):
     devs.bec_worker.planner.move_to(BeamlineState.SAMPLE_ALIGNMENT)
     devs.samcam_auto(AutoEnum.ONCE)
 
+def xtal_snapshot2rse(devs:BeamlineDevices, cfg:BeamlineConfig):
+    common_2rse(devs, cfg)
+
+def xtal_snapshot2se(devs:BeamlineDevices, cfg:BeamlineConfig):
+    sa2se(devs, cfg)
+
+def xtal_snapshot2xrf(devs:BeamlineDevices, cfg:BeamlineConfig):
+    devs.bec_worker.planner.move_to(BeamlineState.FLUX_MEASUREMENT)
+
+def xtal_snapshot2bl(devs: BeamlineDevices, cfg: BeamlineConfig):
+    devs.bec_worker.planner.move_to(BeamlineState.BEAM_VISUALISATION)
+
+def xtal_snapshot2dh(devs: BeamlineDevices, cfg: BeamlineConfig):
+    xtal_snapshot2sa(devs,cfg)
+    sa2dh(devs,cfg)
+
 def dc2rse(devs: BeamlineDevices, cfg: BeamlineConfig):
     devs.samcam_auto(AutoEnum.AUTO)
     common_2rse(devs, cfg)
@@ -167,11 +170,12 @@ def ba2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
     pass
 
 def sa2bl(devs: BeamlineDevices, cfg: BeamlineConfig):
-    devs.scintillator = 20.0
+    devs.bec_worker.planner.move_to(BeamlineState.BEAM_VISUALISATION)
+    devs.samcam_settings.exposure = 0.001
 
 def bl2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
-    devs.scintillator = 20.0
-
+    devs.bec_worker.planner.move_to(BeamlineState.SAMPLE_ALIGNMENT)
+    devs.samcam_auto(AutoEnum.ONCE)
 
 def bl2ba(devs: BeamlineDevices, cfg: BeamlineConfig):
     devs.scintillator = 20.0