csaxs_bec/omny_e2e_tests/sim_flomni_harness.py

"""
Offline harness for the simulated flOMNI hardware stack.

Instantiates the real device classes wired to the protocol simulations
(csaxs_bec.devices.sim) without a running BEC deployment and exercises the code
paths used by the flomni IPython-client scripts:

  1. Galil coarse-stage moves and readback (fgalil protocol, both controllers)
  2. FUPR rotation incl. rotation-angle propagation to the RT controller
  3. drive_axis_to_limit / referencing (#FES) and limit-switch checks
  4. Smaract OSA moves
  5. RT feedback_enable_with_reset + laser tracker (cross-controller choreography)
  6. Fermat-like position scan through RtFlomniFlyer kickoff/complete readout
  7. Gripper sample transfer routine (#MNTMODE, #GRGET incl. confirm handshake)
  8. Simulated cameras (frame generation)

Run with the repo root on PYTHONPATH:  python omny_e2e_tests/sim_flomni_harness.py
"""

from __future__ import annotations

import threading
import time

import numpy as np


class FakeConfiguredDevice:
    """Minimal stand-in for a bec_lib configured device (dm.devices.<name>)."""

    def __init__(self, obj, user_parameter=None):
        self.obj = obj
        self.enabled = True
        self.read_only = False
        self.user_parameter = dict(user_parameter or {})

    def update_user_parameter(self, update: dict):
        self.user_parameter.update(update)


class FakeDeviceContainer(dict):
    def __getattr__(self, name):
        try:
            return self[name]
        except KeyError as exc:
            raise AttributeError(name) from exc


class FakeDeviceManager:
    def __init__(self):
        self.devices = FakeDeviceContainer()

    def add(self, name, obj, user_parameter=None):
        self.devices[name] = FakeConfiguredDevice(obj, user_parameter)
        return obj


def check(condition, message):
    status = "OK  " if condition else "FAIL"
    print(f"[{status}] {message}")
    if not condition:
        raise AssertionError(message)


def main():  # noqa: C901
    from csaxs_bec.devices.sim.sim_cameras import SimIDSCamera, SimWebcamViewer
    from csaxs_bec.devices.sim.sim_galil import SimFlomniGalilMotor, SimFuprGalilMotor
    from csaxs_bec.devices.sim.sim_rt_flomni import SimRtFlomniFlyer, SimRtFlomniMotor
    from csaxs_bec.devices.sim.sim_smaract import SimSmaractMotor

    dm = FakeDeviceManager()
    host = "mpc2844.psi.ch"

    # --- instantiate the simulated endstation -----------------------------------
    fsamx = dm.add(
        "fsamx",
        SimFlomniGalilMotor(
            "E",
            name="fsamx",
            host=host,
            port=8081,
            sign=1,
            limits=[-162, 0],
            device_manager=dm,
            sim_stppermm=2560,
            sim_velocity=50,
            sim_initial_position=-1.14,
        ),
        user_parameter={"in": -1.14},
    )
    fsamy = dm.add(
        "fsamy",
        SimFlomniGalilMotor(
            "F",
            name="fsamy",
            host=host,
            port=8081,
            sign=1,
            limits=[2, 3.8],
            device_manager=dm,
            sim_stppermm=1280000,
            sim_velocity=0.1,
            sim_initial_position=3,
        ),
        user_parameter={"in": 3},
    )
    ftransy = dm.add(
        "ftransy",
        SimFlomniGalilMotor(
            "A",
            name="ftransy",
            host=host,
            port=8081,
            sign=1,
            limits=[-100, 0],
            device_manager=dm,
            sim_stppermm=25600,
            sim_velocity=5,
            sim_initial_position=0,
            sim_analog_inputs={1: -1.6},
        ),
        user_parameter={"sensor_voltage": -1.6},
    )
    ftray = dm.add(
        "ftray",
        SimFlomniGalilMotor(
            "D",
            name="ftray",
            host=host,
            port=8081,
            sign=-1,
            limits=[-200, 0],
            device_manager=dm,
            sim_stppermm=25600,
            sim_velocity=8,
            sim_initial_position=0,
            sim_limit_low_active=True,
        ),
    )
    foptx = dm.add(
        "foptx",
        SimFlomniGalilMotor(
            "B",
            name="foptx",
            host=host,
            port=8082,
            sign=1,
            limits=[-17, -12],
            device_manager=dm,
            sim_stppermm=6400,
            sim_velocity=1,
            sim_initial_position=-13.831,
        ),
    )
    fopty = dm.add(
        "fopty",
        SimFlomniGalilMotor(
            "F",
            name="fopty",
            host=host,
            port=8082,
            sign=1,
            limits=[0, 4],
            device_manager=dm,
            sim_stppermm=128000,
            sim_velocity=0.05,
            sim_initial_position=0.42,
        ),
    )
    foptz = dm.add(
        "foptz",
        SimFlomniGalilMotor(
            "A",
            name="foptz",
            host=host,
            port=8082,
            sign=1,
            limits=[0, 27],
            device_manager=dm,
            sim_stppermm=6400,
            sim_velocity=1,
            sim_initial_position=23,
        ),
    )
    ftrackz = dm.add(
        "ftrackz",
        SimFlomniGalilMotor(
            "G",
            name="ftrackz",
            host=host,
            port=8082,
            sign=1,
            limits=[4.5, 5.5],
            device_manager=dm,
            sim_stppermm=115200,
            sim_velocity=0.02,
            sim_initial_position=5,
        ),
    )
    fsamroy = dm.add(
        "fsamroy",
        SimFuprGalilMotor(
            "A",
            name="fsamroy",
            host=host,
            port=8084,
            sign=-1,
            limits=[-5, 365],
            device_manager=dm,
            sim_initial_position=0,
        ),
    )
    fosax = dm.add(
        "fosax",
        SimSmaractMotor(
            "A",
            name="fosax",
            host=host,
            port=3334,
            sign=-1,
            limits=[10.2, 10.6],
            device_manager=dm,
            sim_initial_position=10.4,
        ),
    )
    rtx = dm.add(
        "rtx",
        SimRtFlomniMotor(
            "A", name="rtx", host=host, port=2222, sign=1, limits=[-200, 200], device_manager=dm
        ),
        user_parameter={"low_signal": 8500, "min_signal": 8000, "rt_pid_voltage": -0.06219},
    )
    rty = dm.add(
        "rty",
        SimRtFlomniMotor(
            "B", name="rty", host=host, port=2222, sign=1, limits=[-100, 100], device_manager=dm
        ),
        user_parameter={"tomo_additional_offsety": 0},
    )
    rtz = dm.add(
        "rtz",
        SimRtFlomniMotor(
            "C", name="rtz", host=host, port=2222, sign=1, limits=[-100, 100], device_manager=dm
        ),
    )
    flyer = dm.add(
        "rt_positions",
        SimRtFlomniFlyer(
            name="rt_positions", host=host, port=2222, device_manager=dm, sim_point_dwell_s=0.005
        ),
    )

    for device in (fsamx, fsamy, ftransy, ftray, foptx, fopty, foptz, ftrackz, fsamroy, fosax):
        device.wait_for_connection()
    for device in (rtx, rty, rtz):
        device.wait_for_connection()
    flyer.controller.on()
    print("--- all simulated devices connected ---")

    # --- 1. galil moves ----------------------------------------------------------
    check(np.isclose(fsamx.readback.get(), -1.14, atol=1e-3), "fsamx boots at -1.14 (in)")
    t0 = time.time()
    fsamx.move(-2.14, wait=True)
    check(np.isclose(fsamx.readback.get(), -2.14, atol=1e-3), "fsamx moved to -2.14")
    fsamx.move(-1.14, wait=True)
    t0 = time.time()
    fsamy.move(3.05, wait=True)
    duration = time.time() - t0
    check(np.isclose(fsamy.readback.get(), 3.05, atol=1e-3), "fsamy moved to 3.05")
    check(duration > 0.3, f"fsamy move at 0.1 mm/s took {duration:.2f} s (constant velocity)")
    move_status = fsamy.move(3.0, wait=False)
    time.sleep(0.25)
    mid_position = fsamy.readback.get()
    while not move_status.done:
        time.sleep(0.05)
    check(
        3.0 < mid_position < 3.05, f"intermediate readback during fsamy move ({mid_position:.4f})"
    )
    check(np.isclose(fsamy.readback.get(), 3.0, atol=1e-3), "fsamy arrived at 3.0")
    fsamy.move(3.05, wait=True)

    sensor = float(ftransy.controller.socket_put_and_receive("MG@AN[1]").strip())
    check(np.isclose(sensor, -1.6), "ftransy gripper sensor voltage reads -1.6")

    # --- 2. FUPR rotation + RT angle propagation ---------------------------------
    fsamroy.move(90, wait=True)
    check(np.isclose(fsamroy.readback.get(), 90, atol=0.1), "fsamroy rotated to 90 deg")
    rt_state = rtx.controller.sock.state
    check(
        np.isclose(rt_state.angle_rad, 90 / 180 * np.pi, atol=1e-3),
        "RT angle updated via FUPR readback",
    )
    fsamroy.move(0, wait=True)

    # --- 3. referencing / limit switches ------------------------------------------
    check(ftray.controller.get_motor_limit_switch("D")[0], "ftray boots with 'IN' limit active")
    foptz.controller.drive_axis_to_limit(0, "reverse")
    check(np.isclose(foptz.readback.get(), 0.0, atol=1e-3), "foptz referenced to 0 at limit")
    check(foptz.controller.get_motor_limit_switch("A")[0], "foptz reverse limit switch active")
    curtain = foptz.controller.fosaz_light_curtain_is_triggered()
    check(curtain, "fosaz light curtain reads triggered")

    # --- 3b. allaxref is a real, writable DMC variable ------------------------------
    con1 = fsamx.controller
    con1.socket_put_confirmed("allaxref=0")
    check(
        con1.socket_put_and_receive("MGallaxref").strip() == "0.0000",
        "allaxref=0 sticks (MG allaxref reads 0)",
    )
    move_refused = False
    try:
        fsamx.move(-1.2, wait=True)
    except Exception:
        move_refused = True
    check(move_refused, "move refused while allaxref=0")
    con1.socket_put_confirmed("allaxref=1")
    check(con1.all_axes_referenced(), "allaxref=1 restores referenced state")
    fsamx.move(-1.14, wait=True)

    con1.socket_put_and_receive("AB")
    check(
        not con1.is_thread_active(5) and not fsamx.controller.is_axis_moving("E", 4),
        "AB aborts motion and program threads",
    )

    # --- 4. smaract ---------------------------------------------------------------
    fosax.move(10.5, wait=True)
    check(np.isclose(fosax.readback.get(), 10.5, atol=1e-4), "fosax moved to 10.5")
    check(fosax.controller.axis_is_referenced(0), "fosax referenced")

    # --- 5. RT feedback + tracker --------------------------------------------------
    rtx.controller.feedback_enable_with_reset()
    check(rtx.controller.feedback_is_running(), "interferometer feedback running")
    check(rtx.controller.laser_tracker_check_enabled(), "laser tracker enabled")
    check(dm.devices["fsamx"].read_only, "fsamx set to read-only with feedback on")
    signal_status = rtx.controller.laser_tracker_check_signalstrength(verbose=False)
    check(signal_status == "ok", f"tracker signal strength '{signal_status}'")

    rtx.move(5.0, wait=True)
    check(np.isclose(rtx.readback.get(), 5.0, atol=1e-3), "rtx moved to 5.0 with feedback")
    rtx.controller.move_samx_to_scan_region(5.0)

    # --- 6. fermat-like scan through the flyer -------------------------------------
    num_points = 400
    positions = [[5.0 + 3 * np.sin(i * 0.3), 2 * np.cos(i * 0.3), 0.0] for i in range(num_points)]
    rtx.controller.clear_trajectory_generator()
    rtx.controller.add_pos_to_scan(positions)
    flyer.stage()
    status = flyer.kickoff()
    status.wait(timeout=10)
    complete_status = flyer.complete()
    deadline = time.time() + 30
    while not complete_status.done and time.time() < deadline:
        time.sleep(0.1)
    check(complete_status.done, "flyer completed the scan readout")
    data = flyer.data.get()
    check(rtx.controller.get_scan_status()[0] == 0, "scan mode back to 0 after scan")
    print(f"       last data row: {data}")

    # --- 7. gripper transfer routine -----------------------------------------------
    con = ftransy.controller
    con.socket_put_confirmed("XQ#MNTMODE")
    time.sleep(0.1)
    check(bool(float(con.socket_put_and_receive("MG mntmod").strip())), "controller in mount mode")
    con.socket_put_confirmed("getaprch=-17.5")
    con.socket_put_confirmed("XQ#GRGET,3")

    def confirmer():
        for _ in range(200):
            confirm = int(float(con.socket_put_and_receive("MG confirm").strip()))
            if confirm == -1:
                con.socket_put_confirmed("confirm=1")
            in_progress = bool(float(con.socket_put_and_receive("MG mntprgs").strip()))
            if not in_progress:
                return
            time.sleep(0.1)

    confirm_thread = threading.Thread(target=confirmer)
    confirm_thread.start()
    deadline = time.time() + 30
    while time.time() < deadline:
        if not bool(float(con.socket_put_and_receive("MG mntprgs").strip())):
            break
        time.sleep(0.2)
    confirm_thread.join()
    check(
        not bool(float(con.socket_put_and_receive("MG mntprgs").strip())),
        "gripper #GRGET routine finished (confirm handshakes served)",
    )
    con.socket_put_confirmed("XQ#POSMODE")
    time.sleep(0.1)
    check(
        not bool(float(con.socket_put_and_receive("MG mntmod").strip())),
        "controller back in positioning mode",
    )

    # --- 8. cameras -----------------------------------------------------------------
    webcam = SimWebcamViewer(url="http://sim/video", name="cam_flomni_gripper")
    ids_cam = SimIDSCamera(name="cam_xeye", camera_id=11, force_monochrome=True, sim_noise_std=4.0)
    frame = ids_cam.cam.get_image_data()
    check(frame.ndim == 2 and frame.max() > 200, f"IDS sim frame {frame.shape}, max {frame.max()}")
    check(webcam._sim_frame.shape == (480, 640, 3), "webcam sim frame 480x640 RGB")

    # --- 9. sample storage -----------------------------------------------------------
    from csaxs_bec.devices.sim.sim_sample_storage import SimFlomniSampleStorage

    samples = SimFlomniSampleStorage(name="flomni_samples")
    check(
        samples.is_sample_slot_used(0) and samples.is_sample_slot_used(3),
        "sample storage boots with samples in slots 0-3",
    )
    check(not samples.is_sample_in_gripper(), "gripper boots empty")
    samples.set_sample_in_gripper(samples.get_sample_name(1))
    samples.unset_sample_slot(1)
    check(
        samples.is_sample_in_gripper() and not samples.is_sample_slot_used(1),
        "slot 1 -> gripper bookkeeping",
    )

    # --- 9b. sample owner (packed into the same DESC field as name) ------------------
    samples.set_sample_slot(2, "owner_test_sample", owner="mholler")
    check(
        samples.get_sample_name(2) == "owner_test_sample", "get_sample_name returns plain name only"
    )
    check(samples.get_sample_owner(2) == "mholler", "get_sample_owner returns the packed owner")
    check(
        samples.get_sample_name_and_owner(2) == ("owner_test_sample", "mholler"),
        "get_sample_name_and_owner returns (name, owner)",
    )
    samples.unset_sample_slot(2)
    check(
        str(getattr(samples.sample_names, "sample2").get()) == "-",
        "unset_sample_slot writes bare '-', not '- | <owner>'",
    )
    # gripper move (slot -> gripper) must preserve owner, not just name. The
    # gripper is a separate Cpt (sample_in_gripper_name), not part of the
    # sample_names slot dict, so unpack it directly rather than via
    # get_sample_name_and_owner (which only covers slots 0-20).
    from csaxs_bec.devices.omny.sample_desc_codec import unpack_desc

    samples.set_sample_slot(4, "transfer_test_sample", owner="e12345")
    name, owner = samples.get_sample_name_and_owner(4)
    samples.set_sample_in_gripper(name, owner=owner)
    samples.unset_sample_slot(4)
    check(
        unpack_desc(str(samples.sample_in_gripper_name.get())) == ("transfer_test_sample", "e12345"),
        "owner survives a slot -> gripper move",
    )
    samples.unset_sample_in_gripper()

    frame_a = ids_cam.cam.get_image_data()
    frame_b = ids_cam.cam.get_image_data()
    check(not np.array_equal(frame_a, frame_b), "IDS frames are live (noise explicitly enabled)")
    static_cam = SimIDSCamera(name="cam_xeye_static", camera_id=12, force_monochrome=True)
    check(
        np.array_equal(static_cam.cam.get_image_data(), static_cam.cam.get_image_data()),
        "IDS frames are static by default (noise disabled)",
    )

    # --- feedback disable puts stages back to writable -------------------------------
    rtx.controller.feedback_disable()
    check(not dm.devices["fsamx"].read_only, "fsamx writable again after feedback_disable")

    print("\nAll harness checks passed.")


if __name__ == "__main__":
    main()