AareDAQ/src/aare/devices/bec_worker.py

import sys
import time
from enum import Enum

from aarecommon.config.beamline import cfg_get, mx_beamline
from aarecommon.config.logger import setup_logger
from aarecommon.config.logger_events import log_timing
from aarecommon.errors.exception_handler import BECCommunicationError
from aarecommon.models.beamline import MXBeamline
from aarecommon.models.models import BeamlineStateEnum
from bec_ipython_client import BECIPythonClient
from bec_ipython_client.signals import OperationMode
from bec_lib.device import RPCError, ScanRequestError
from bec_lib.service_config import ServiceConfig

from aare.beamline_dispatch.beamline_dispatch import get_beamline_dispatch

logger = setup_logger("aareDAQ")


class DetectorCoverEnum(str, Enum):
    """Enum for the detector cover position
    position devices can take string or number to move
    Currently dictated string in DAQ for consitency"""

    OPEN = "open"  # 2
    CLOSED = "closed"  # 1


class BrightnessEnum(str, Enum):
    """Enum for the backlight brightness
    position devices can take string or number to move
    """

    ON = "on"
    OFF = "off"


class BeamlineState(str, Enum):
    ROBOT_SAMPLE_EXCHANGE = "robot_sample_exchange"
    SAMPLE_ALIGNMENT = "sample_alignment"
    DATA_COLLECTION = "data_collection"
    DC_XRF = "DC_XRF"
    MANUAL_SAMPLE_EXCHANGE = "manual_sample_exchange"
    BEAM_VISUALISATION = "beam_visualisation"
    FLUX_MEASUREMENT = "flux_measurement"
    BEAMSTOP_ALIGNMENT = "beamstop_alignment"
    MAINTENANCE = "maintenance"
    XTAL_SNAPSHOT = "xtal_snapshot"


def _bec_state_to_aare_state(bec_state: BeamlineState) -> BeamlineStateEnum:
    map = {
        BeamlineState.BEAMSTOP_ALIGNMENT: BeamlineStateEnum.BeamstopAlignment,
        BeamlineState.SAMPLE_ALIGNMENT: BeamlineStateEnum.SampleAlignment,
        BeamlineState.DATA_COLLECTION: BeamlineStateEnum.DataCollection,
        BeamlineState.DC_XRF: BeamlineStateEnum.XrayFluorescence,
        BeamlineState.MANUAL_SAMPLE_EXCHANGE: BeamlineStateEnum.SampleExchange,
        BeamlineState.BEAM_VISUALISATION: BeamlineStateEnum.BeamLocation,
        BeamlineState.FLUX_MEASUREMENT: BeamlineStateEnum.FluxMeasurement,
        BeamlineState.BEAMSTOP_ALIGNMENT: BeamlineStateEnum.BeamstopAlignment,
        BeamlineState.MAINTENANCE: BeamlineStateEnum.Maintenance,
        BeamlineState.XTAL_SNAPSHOT: BeamlineStateEnum.XtalSnapshot,
    }
    if bec_state in map:
        return map[bec_state]
    return BeamlineStateEnum.Maintenance


class BECClientWorker:
    def __init__(self, beamline: MXBeamline, name: str = "default"):
        logger.debug(f"initialising BEC worker for {beamline}")
        BEAMLINE = beamline.value.lower()
        self.beamline = beamline
        if self.beamline is MXBeamline.X06DA:
            self._beamline_name = "pxiii"
        elif self.beamline is MXBeamline.X10SA:
            self._beamline_name = "pxii"
        elif self.beamline is MXBeamline.X06DA:
            self._beamline_name = "pxi"
        elif self.beamline is MXBeamline.SIMULATED:
            self._beamline_name = "SIMULATED"
        else:
            raise ValueError(f"Unknown beamline: {beamline}")

        self.dispatch = get_beamline_dispatch()

        if self.beamline is MXBeamline.SIMULATED:
            self.simulated = True
            host = "localhost"
        else:
            self.simulated = False
            host = cfg_get("daq.hardware.bec_url", f"{BEAMLINE}-bec-001.psi.ch")
        service_config = ServiceConfig(redis={"host": host, "port": 6379})
        service_config.config["log_writer"]["base_path"] = "/tmp/logs"
        try:
            self.client = BECIPythonClient(config=service_config, mode=OperationMode.Procedure)
            self.client.start()
            self.dev = self.client.device_manager.devices
            self.scans = self.client.scans
            self.macros = self.dispatch.bec_macros
            self._set_scilog_tags()
            self._init_beamline_environment()
        except Exception:
            logger.exception("Error initialising BEC devices")
            self.client.shutdown()
            sys.exit(1)
        logger.debug(f"simulated is {self.simulated}")

    def shutdown(self):
        self.client.shutdown()

    def _init_beamline_environment(self):
        try:
            self.position_devices, self.planner = self.macros.init_beamline_environment()
            self._backlight_brightness = self.position_devices["bl_bright"]
            self._frontlight_brightness = self.position_devices["fl_bright"]
            self._zoom = self.dev.scam_zoom
            self._ring_current = self.dev.sls_current
        except Exception as e:
            logger.error(f"Error initialising planner and position devices: {e}")
            self.position_devices = None
            self.planner = None
            self._backlight_brightness = None
            self._frontlight_brightness = None
            try:
                self._zoom = self.dev.scam_zoom
                self._ring_current = self.dev.sls_current
            except Exception as e:
                logger.exception("Error initialising zoom and ring_current")
                self._zoom = None
                self.ring_current = None
            raise RuntimeError(f"Error initialising BEC devices: {e}") from e

    def read_current_state(self) -> BeamlineStateEnum:
        if self.planner is None:
            return BeamlineStateEnum.Maintenance
        matching_states = self.planner.current_state()
        if matching_states is None or len(matching_states) > 1:
            return BeamlineStateEnum.Maintenance
        return _bec_state_to_aare_state(matching_states[0])

    def _bec_error(
        self, exc: Exception, *, operation: str, tags: list[str] | None = None
    ) -> BECCommunicationError:
        """Report a failed BEC operation and build the error; callers ``raise ... from`` it."""
        message = f"BEC operation '{operation}' failed: {type(exc).__name__}: {exc}"
        # logger.exception(message)
        if tags is None:
            tags = ["error", "bec"]
        else:
            tags += ["error", "bec"]
        try:
            bec_alarms = self.client.alarms()
            if bec_alarms:
                alarms = [alarm for alarm in bec_alarms]
                logger.debug(f"alarms: {alarms}")
                last_alarm = self.client.show_last_alarm()
                message += f"\n\nAlarm: {last_alarm}\n\n"
                logger.error(f"last_alarm: {last_alarm}")
        except Exception:
            logger.exception("Couldn't raise BEC alarms")

        try:
            self.scilog_msg(
                message=message,
                error=True,
                error_message=f"Error during '{operation}': {exc}",
                tags=tags,
            )
        except Exception:
            logger.exception("Error sending scilog message")

        if isinstance(exc, AssertionError):
            return BECCommunicationError(
                f"BEC internal assertion failed during '{operation}'",
                operation=operation,
                exception=exc,
            )

        return BECCommunicationError(message, operation=operation, exception=exc)

    def _set_scilog_tags(self, tags: list[str] | None = None):
        try:
            if tags:
                self.client.messaging.scilog.set_default_tags(tags)
            else:
                self.client.messaging.scilog.set_default_tags(["AareDAQ"])
        except Exception as e:
            logger.exception("Error setting scilog tags")
            raise self._bec_error(e, operation="set scilog tags") from e

    def scilog_msg(
        self,
        message: str,
        error: bool = False,
        warning: bool = False,
        error_message: str | None = None,
        attachments=None,
        bold: bool = False,
        italic: bool = False,
        color: str | None = None,
        additonal_text: list[str] | None = None,
        tags: list[str] | None = None,
    ):
        if color and color not in ["red", "green", "yellow", "blue", "pink"]:
            logger.warning("specified color not in allowed list,using default")
            color = None
        try:
            msg = self.client.messaging.scilog.new()
        except Exception as e:
            raise self._bec_error(e, operation="scilog_msg") from e
        try:
            msg.add_text(message, bold=bold, italic=italic, color=color)
            if error:
                msg.add_text(error_message, bold=True, color="red")
            elif warning:
                msg.add_text(error_message, bold=True, color="yellow")
        except Exception as e:
            logger.exception("Error adding text")
            msg.add_text(f"Error adding text: {e}")
        try:
            if attachments:
                for attachment in attachments:
                    msg.add_attachment(attachment)
        except Exception as e:
            logger.exception("Error adding attachment")
            msg.add_text(f"Error adding attachment: {e}")
        try:
            if additonal_text:
                for text in additonal_text:
                    msg.add_text(text)
        except Exception as e:
            logger.exception("Error adding additional text")
            msg.add_text(f"Error adding additional text: {e}")
        try:
            if tags:
                msg.add_tags(tags)
        except Exception as e:
            logger.exception("Error setting scilog tags")
            msg.add_text(f"Error setting scilog tags: {e}")
        try:
            msg.send()
        except Exception:
            logger.exception("Error sending scilog message")

    @log_timing(logger, "BEC move_to")
    def move_to(self, state: BeamlineState):
        start = time.perf_counter()
        logger.debug(f"simulated is {self.simulated}")
        logger.info(f"BEC move_to requested: {state.value}")
        if self.simulated:
            logger.debug(f"Simulating move to {state.value}")
            return True
        try:
            self.planner.move_to(state)
            if self.planner.is_state(state):
                logger.info(
                    f"BEC move_to completed: {state.value} in {time.perf_counter() - start:.2f}s"
                )
                return True
            else:
                logger.warning(
                    f"Failed to move to {state.value} in {time.perf_counter() - start:.2f}s"
                )
                return False
        except Exception as e:
            raise self._bec_error(e, operation=f"planner.move_to:{state.value}") from e

    def is_state(self, state: BeamlineState):
        if self.simulated:
            logger.debug(f"Simulating check_beamline_state: {state.value}")
            return True
        return self.planner.is_state(state)

    def list_position_devices(self):
        """List the position devices available for the BEC worker"""
        if self.simulated or self.position_devices is None:
            return []
        return list(self.position_devices.keys())

    def show_all_devices(self):
        return self.dev.show_all

    def list_all_user_macros(self) -> list[str] | None:
        if self.simulated:
            logger.debug("Simulating list_all_user_macros")
            return []
        try:
            self._list_all_macros()
            raw_macros = [name for name, _ in self.client.macros._update_handler.macros.items()]
            if raw_macros is None:
                logger.warning("BEC returned no user macros; treating as empty list")
                return []
            return [str(macro) for macro in raw_macros]
        except Exception as e:
            raise self._bec_error(e, operation="list_all_user_macros") from e

    def _list_all_macros(self):
        result = self.client.macros.list_user_macros()
        if result is None:
            return []
        return result

    def load_user_macros(self):
        if self.simulated:
            logger.debug("Simulating load_user_macros")
            return None
        try:
            return self._load_user_macros()
        except Exception as e:
            raise self._bec_error(e, operation="load_user_macros") from e

    def _load_user_macros(self):
        result = self.macros.load_all_user_macros()
        if result is None:
            logger.warning("BEC load_all_user_macros returned None")
        return result

    def reinitialise_planner_and_position_devices(self, method: str = "auto"):
        """
        Reinitialise BEC position devices and planner.

        Args:
            method:
                "auto" - use the beamline-specific default.
                "beamline" - force init_beamline_environment().
                "sample" - force init_se_devices() + planner creation.

        Returns:
            List of position device names after reinitialisation.
        """
        if self.simulated:
            logger.debug(f"Simulating reinitialise_planner_and_position_devices(method={method})")
            return []
        try:
            self.client.config.update_session_with_file(
                f"/sls/{self.beamline}/config/bec/production/{self._beamline_name}_bec/{self._beamline_name}_bec/device_configs/{self._beamline_name}-devices.yaml"
            )
            self._init_beamline_environment()
            logger.info(f"Reinitialised BEC planner and position devices using method={method}")
            return self.list_position_devices()
        except Exception as e:
            raise self._bec_error(
                e, operation=f"reinitialise_planner_and_position_devices:{method}"
            ) from e

    def shutdown_client(self):
        self.client.shutdown()

    def mono_pitch_scan_runner(self, plot: bool = False):
        try:
            self.macros.mono_pitch_scan(plot)
        except Exception as e:
            raise self._bec_error(e, operation="mono_pitch_scan", tags=["mono_pitch_scan"]) from e
        if self.beamline is MXBeamline.X10SA:
            additional_text = [f"New dcm_pitch position: {self.dev.dcm_pitch.position:5f}"]
        elif self.beamline is MXBeamline.X06DA:
            additional_text = [f"New dcm_theta2 position: {self.dev.dccm_theta2.position:5f}"]
        else:
            additional_text = ["What beamline did you run this on???"]
        self.scilog_msg(
            message="Mono pitch scan completed",
            bold=True,
            color="green",
            tags=["mono_pitch_scan"],
            additonal_text=additional_text,
        )

    def check_current_energy(self):
        """Get the current energy from the BEC in eV"""
        energy_ev = self.macros.get_current_energy()
        energy_kev = energy_ev / 1000
        return energy_kev

    def change_energy(self, value: float, plot: bool = False):
        current_energy = self.check_current_energy()
        logger.info(f"Current energy: {current_energy:.1f} eV")
        logger.info(f"Change energy requested: from {current_energy:.1f} to {value:.1f} eV")
        try:
            self.macros.bl_energy(value, move_gap=False, mono_scan=True, plot=plot)
        except Exception as e:
            raise self._bec_error(
                e,
                operation=f"Requested energy change from:{current_energy:.1f} to {value} eV",
                tags=["energy_change"],
            ) from e

        if abs(value - self.check_current_energy()) > 1:
            logger.warning(
                f"Energy change may have failed, current energy: {self.check_current_energy()} eV"
            )
        if self.beamline is MXBeamline.X10SA:
            additonal_text = [
                f"New dcm_bragg position: {self.dev.dcm_bragg.position:4g} mrad",
                f"New dcm_pitch position: {self.dev.dcm_pitch.position:4g} ",
                f"Previous energy: {current_energy:.1f} eV ",
                f"Requested energy: {value:.1f} eV ",
                f"New current energy: {self.check_current_energy():.1f} eV",
            ]
        else:
            additonal_text = [
                f"New dccm_theta1 position: {self.dev.dccm_theta1.position:4g} mrad",
                f"New dccm_theta2 position: {self.dev.dccm_theta2.position:4g} mrad",
                f"Previous energy: {current_energy:.1f} eV ",
                f"Requested energy: {value:.1f} eV ",
                f"New current energy: {self.check_current_energy():.1f} eV",
            ]
        self.scilog_msg(
            message=f"Moved from {current_energy:.1f} eV to {value:.1f} eV",
            bold=True,
            color="green",
            tags=["energy_change"],
            additonal_text=additonal_text,
        )

    def get_det_z(self):
        try:
            return self.dev.det_z.position
        except Exception as e:
            raise self._bec_error(e, operation="get_det_z", tags=["det_z"]) from e

    def det_z(self, value: float, timeout: int | None = None):
        """timeout is None or integer in s"""
        try:
            status = self.scans.mv(self.dev.det_z, value, relative=False)
            if timeout:
                status.wait(timeout=timeout)
            return status
        except Exception as e:
            raise self._bec_error(e, operation=f"scans.mv:det_z:{value}", tags=["det_z"]) from e

    def get_det_y(self):
        try:
            return self.dev.det_y.position
        except Exception as e:
            raise self._bec_error(e, operation="get_det_z", tags=["det_z"]) from e

    def det_y(self, value: float, timeout: int | None = None):
        """timeout is None or integer in s"""
        try:
            status = self.scans.mv(self.dev.det_y, value, relative=False)
            if timeout:
                status.wait(timeout=timeout)
            return status
        except Exception as e:
            raise self._bec_error(e, operation=f"scans.mv:det_y:{value}", tags=["det_z"]) from e

    @property
    def backlight_brightness(self) -> BrightnessEnum:
        """returns backlight brightness as an Enum: 'off' or 'on', can also be a value...
        How to handle"""
        try:
            return BrightnessEnum(self._backlight_brightness.actual)
        except Exception as e:
            raise self._bec_error(
                e,
                operation="backlight brightness, could not get backlight brightness",
                tags=["backlight"],
            ) from e
            raise

    @backlight_brightness.setter
    def backlight_brightness(self, value: int | str):
        """Set the backlight brightness to the specified value"""
        if self.simulated:
            return
        try:
            self._backlight_brightness.move(value)
        except Exception as e:
            raise self._bec_error(
                e, operation=f"backlight_brightness:{value}", tags=["backlight"]
            ) from e
            raise

    def get_backlight_pos(self) -> BrightnessEnum:
        """Returns the current backlight brightness position"""
        return BrightnessEnum(self._backlight_brightness.pos)

    def backlight_toggle(self):
        """Turn the backlight on or off"""
        if self.simulated:
            return
        try:
            current = self.get_backlight_pos()
            if current is BrightnessEnum.ON:
                target = BrightnessEnum.OFF
            else:
                target = BrightnessEnum.ON
            self.backlight_brightness = target
        except Exception as e:
            raise self._bec_error(
                e,
                operation="backlight toggle, could not change backlight on/off ",
                tags=["backlight"],
            ) from e

    def save_current_bs_pos(self):
        self.macros.save_current_position(self.dev.bs_y, "in")
        self.macros.save_current_position(self.dev.bs_x, "in")

    def save_current_collimator_pos(self):
        self.macros.save_current_position(self.dev.coll_x, "in")
        self.macros.save_current_position(self.dev.coll_y, "in")

    def save_current_aerotech_position(self):
        self.macros.save_current_position(self.dev.aerotech, "in", axis="x")
        self.macros.save_current_position(self.dev.aerotech, "work", axis="y")
        self.macros.save_current_position(self.dev.aerotech, "work", axis="z")

        self.save_config_and_reload_devices()

    def save_config_and_reload_devices(self):
        self.position_devices, self.planner = self.macros.save_and_reload()

    @property
    def zoom(self):
        return self._zoom.position

    @zoom.setter
    def zoom(self, value: float):
        self.scans.umv(self._zoom, value, relative=False)

    @property
    def ring_current(self) -> float:
        try:
            return self._ring_current.get()
        except (RPCError, ScanRequestError):
            return 0


if __name__ == "__main__":
    import time

    print(time.ctime(), " starting BEC Client")
    beamline = mx_beamline()
    try:
        client = BECClientWorker(beamline)
    except Exception:
        logger.exception("Failed to start BEC client")
        try:
            client.shutdown_client()
        except Exception:
            logger.exception("Failed to shut the BEC client down cleanly")
            sys.exit(1)
    # print(client.get_det_cov(actual=True))
    # print(client.is_state(BeamlineState.ROBOT_SAMPLE_EXCHANGE))
    try:
        print("startting backlight brightness test")
        print("initial value")
        # print(client.client.show_last_alarm())
        # print(client._raise_bec_error(exc=Exception("test"), operation="test", tags=["test"]))
        print(client.ring_current)
        for i in range(10):
            print(client._ring_current.get())
            time.sleep(1)
        # print('setting to 5')
        # client.backlight_brightness = 5
        # print(client.get_backlight_pos())
        # print(client.backlight_brightness)
        # print('setting to 1.2')
        # client.backlight_brightness = 1.2
        # print(client.get_backlight_pos())
        # print(client.backlight_brightness)
        # print('setting to ON')
        # client.backlight_brightness = BrightnessEnum.ON
        # print(client.get_backlight_pos())
        # print(client.backlight_brightness)
        # print('setting to 1.0')
        # client.backlight_brightness = 1.0
        # print(client.get_backlight_pos())
        # print(client.backlight_brightness)
        # print('setting to 0.0')
        # client.backlight_brightness = 0.0
        # print(client.get_backlight_pos())
        # print(client.backlight_brightness)
        # print('setting to OFF')
        # client.backlight_brightness = BrightnessEnum.OFF
        # client.backlight_toggle()
        # print(client.get_backlight_pos())
        # print(client.backlight_brightness)
        # client.scilog_msg("Testing scilog messages with color = yellow and italic", italic=True,
        #                  color="green", warning=False)
    except Exception:
        client.shutdown_client()
        logger.exception("BEC client smoke test failed")

    # try:
    #     det_value = 980
    #     print(f"moving detector to vale:{det_value}")
    #     start = time.perf_counter()
    #     # print(client.get_det_z())
    #     # client.det_z(value=det_value, timeout=10)
    #     # print(time.perf_counter() - start)
    #     # print(client.get_det_z)
    #     # det_value = 985
    #     # print(f"moving detector to vale:{det_value}")
    #     status = client.det_z(value=det_value)
    #     status.wait(timeout=5)
    #     print(status)
    #     print(client.get_det_z())
    #     #client.mono_pitch_scan_runner()
    #     #client.change_energy(12000)
    #
    # except Exception as e:
    #     print(f"Error: {e}")
    # # client.planner.current_state()
    # client.planner.move_to(BeamlineState.SAMPLE_ALIGNMENT)
    # time.sleep(2.0)
    # client.planner.current_state()
    # client.move_to(BeamlineState.SAMPLE_ALIGNMENT)
    # time.sleep(2.0)
    # client.planner.current_state()

    # client.load_user_macros()
    # client.macros.mono_pitch_scan(False)
    # status = client.run_macro("planner.current_state", queue="default")
    # client.planner.move_to('manual_sample_exchange')
    # print(status)
    # status.wait()
    # print(status)
    # try:
    # status=client.proc.run_macro("mono_pitch_scan", queue="test")
    # planner.move_to('manual_sample_exchange')

    # print(status)
    # status.wait()
    # status.cancel()
    # print(status)
    # client.client.macros.mono_pictch_scan(False)
    # try:
    #     # a=client.a2e_runner(160, "iln")
    #     # print(a)
    #     # print(convert_from_energy(12))
    #     # energy = get_current_energy()
    #     # pos = get_dcm_motors_positions(energy)
    #     # print(energy, pos)
    #     print(bs_z_policy(15.0))
    #     #client.scans.umv(client.dev.xeye_x, 0, relative=False)
    #     #client.mono_pitch_scan_runner()
    #     #client.mono_pitch_scan_runner()
    #     #b=client.run_macro_blocked("a2e", 160, "iln", queue="test")
    #     #b = client.run_macro_blocked("mono_pitch_scan", False, queue="default")
    #
    #     #status = client.mono_pitch_scan_runner
    #     #print(status)
    #
    #     #client.rse2sa()
    #     #time.sleep(10.0)
    #     #time.sleep(10.0)
    #     #print(status)
    #     #client.common2rse()
    #     #print(status)
    # except GuardViolation as e:
    #     print(f"GuardViolation: {e}")
    # except RuntimeError as e:
    #     print(f"RuntimeError: {e}")
    client.shutdown_client()