"""
Generate simulated device configs for LamNI and OMNY from the real ptycho configs.

Mirrors ptycho_lamni.yaml / ptycho_omny.yaml, swapping deviceClass to the simulated
counterparts and injecting sim_* seed parameters (resolutions from lamni.dmc, boot
positions from the 'in' user parameters where available). Devices without a simulated
counterpart are dropped and reported.

Run from the repo root:  python bin/generate_simulated_configs.py
"""

from __future__ import annotations

import sys
from pathlib import Path

import yaml

REPO = Path(__file__).resolve().parents[1]
CONFIGS = REPO / "csaxs_bec" / "device_configs"

CLASS_MAP = {
    "csaxs_bec.devices.omny.galil.lgalil_ophyd.LamniGalilMotor": "csaxs_bec.devices.sim.sim_lamni.SimLamniGalilMotor",
    "csaxs_bec.devices.omny.rt.rt_lamni_ophyd.RtLamniMotor": "csaxs_bec.devices.sim.sim_lamni.SimRtLamniMotor",
    "csaxs_bec.devices.omny.rt.rt_lamni_ophyd.RtLamniFlyer": "csaxs_bec.devices.sim.sim_lamni.SimRtLamniFlyer",
    "csaxs_bec.devices.omny.galil.ogalil_ophyd.OMNYGalilMotor": "csaxs_bec.devices.sim.sim_omny.SimOMNYGalilMotor",
    "csaxs_bec.devices.omny.rt.rt_omny_ophyd.RtOMNYMotor": "csaxs_bec.devices.sim.sim_omny.SimRtOMNYMotor",
    "csaxs_bec.devices.omny.rt.rt_omny_ophyd.RtOMNYFlyer": "csaxs_bec.devices.sim.sim_omny.SimRtOMNYFlyer",
    "csaxs_bec.devices.omny.omny_sample_storage.OMNYSampleStorage": "csaxs_bec.devices.sim.sim_omny.SimOMNYSampleStorage",
    "csaxs_bec.devices.omny.omny_dewar.OMNYDewar": "csaxs_bec.devices.sim.sim_omny.SimOMNYDewar",
    "csaxs_bec.devices.omny.omny_temperatures.OMNYTemperatures": "csaxs_bec.devices.sim.sim_omny.SimOMNYTemperatures",
    "csaxs_bec.devices.smaract.smaract_ophyd.SmaractMotor": "csaxs_bec.devices.sim.sim_smaract.SimSmaractMotor",
    "csaxs_bec.devices.ids_cameras.ids_camera.IDSCamera": "csaxs_bec.devices.sim.sim_cameras.SimIDSCamera",
    "csaxs_bec.devices.omny.webcam_viewer.WebcamViewer": "csaxs_bec.devices.sim.sim_cameras.SimWebcamViewer",
    "csaxs_bec.devices.omny.flomni_sample_storage.FlomniSampleStorage": "csaxs_bec.devices.sim.sim_sample_storage.SimFlomniSampleStorage",
    "ophyd_devices.ComputedSignal": "ophyd_devices.ComputedSignal",
}

# lamni.dmc #INIT: stppermm / encpermm per axis index (A=0 .. H=7)
LAMNI_STPPERMM = {
    0: 12800,
    1: 12800,
    2: 50154.32099,
    3: 12800,
    4: 12800,
    5: 38400,
    6: 12800,
    7: 12800,
}
LAMNI_ENCPERMM = {0: 2000, 1: 2000, 2: 36000, 3: 10000, 4: 10000, 5: 30000}
LAMNI_ENCODER_AXES = set(range(6))
# lamni.dmc #INIT axspeed table (units/s)
LAMNI_VELOCITY = {0: 3, 1: 3, 2: 12, 3: 1, 4: 0.5, 5: 0.5, 6: 3, 7: 3}

SIM_BEAMLINE_DEVICES = """
ddg1:
  description: Simulated main delay generator for triggering
  deviceClass: csaxs_bec.devices.sim.simulated_beamline_devices.SimulatedDDG1
  enabled: true
  deviceConfig:
    prefix: 'X12SA-CPCL-DDG1:'
  onFailure: raise
  readOnly: false
  readoutPriority: baseline
  softwareTrigger: true
fsh:
  description: Simulated fast shutter manual control and readback
  deviceClass: csaxs_bec.devices.sim.simulated_beamline_devices.cSAXSSimulatedFastShutter
  deviceConfig:
    prefix: 'X12SA-ES1-TTL:'
  onFailure: raise
  enabled: true
  readoutPriority: monitored
"""


def axis_index(axis_id: str) -> int:
    return ord(axis_id.lower()) - ord("a")


def transform(src_name: str, dst_dir: str, dst_name: str, setup: str) -> None:
    src = CONFIGS / src_name
    config = yaml.safe_load(src.read_text())
    out, dropped = {}, []
    for name, dev in config.items():
        if not isinstance(dev, dict) or "deviceClass" not in dev:
            dropped.append(name)
            continue
        dev_class = dev["deviceClass"].strip()
        if dev_class not in CLASS_MAP:
            dropped.append(f"{name} ({dev_class.rsplit('.', 1)[-1]})")
            continue
        dev = dict(dev)
        dev["deviceClass"] = CLASS_MAP[dev_class]
        dev_config = dict(dev.get("deviceConfig") or {})
        user_params = dev.get("userParameter") or {}

        if "SimLamniGalilMotor" in dev["deviceClass"]:
            axis = axis_index(dev_config["axis_Id"])
            dev_config["sim_stppermm"] = LAMNI_STPPERMM.get(axis, 12800)
            if axis in LAMNI_ENCODER_AXES:
                dev_config["sim_encpermm"] = LAMNI_ENCPERMM[axis]
            dev_config.setdefault("sim_velocity", LAMNI_VELOCITY.get(axis, 1))
        if "GalilMotor" in dev["deviceClass"] or "SmaractMotor" in dev["deviceClass"]:
            # boot at the aligned position: 'center' (LamNI) or 'in' user parameter
            initial = user_params.get("center", user_params.get("in", 0))
            limits = dev_config.get("limits")
            if isinstance(limits, list) and len(limits) == 2 and None not in limits:
                low, high = sorted(float(limit) for limit in limits)
                # low == high means 'no limit check' in BEC; do not clamp then
                if low < high and not low <= float(initial) <= high:
                    initial = round((low + high) / 2, 4)
            dev_config["sim_initial_position"] = initial
        if "SimRtOMNYFlyer" in dev["deviceClass"] or "SimRtLamniFlyer" in dev["deviceClass"]:
            dev_config.setdefault("sim_point_dwell_s", 0.02)

        dev["deviceConfig"] = dev_config
        out[name] = dev

    # LamNI: couple the rt interferometer readback to the coarse stages so that
    # lamni_new_scan_center_interferometer converges (see SimRtLamniState)
    if setup == "LamNI" and "lsamx" in out and "lsamy" in out:
        coupling = {
            "galil_port": out["lsamx"]["deviceConfig"]["port"],
            "x_axis": axis_index(out["lsamx"]["deviceConfig"]["axis_Id"]),
            "x_sign": out["lsamx"]["deviceConfig"].get("sign", 1),
            "x_center": (config["lsamx"].get("userParameter") or {}).get("center", 0),
            "y_axis": axis_index(out["lsamy"]["deviceConfig"]["axis_Id"]),
            "y_sign": out["lsamy"]["deviceConfig"].get("sign", 1),
            "y_center": (config["lsamy"].get("userParameter") or {}).get("center", 0),
        }
        for rt_name in ("rtx", "rty", "rt_positions"):
            if rt_name in out:
                out[rt_name]["deviceConfig"]["sim_coarse_coupling"] = coupling

    header = (
        f"############################################################\n"
        f"# Simulated {setup} endstation (generated by bin/generate_simulated_configs.py)\n"
        f"# Real device classes run against protocol-level simulations.\n"
        f"# Dropped (no simulated counterpart): {', '.join(dropped) if dropped else 'none'}\n"
        f"############################################################\n\n"
    )
    body = yaml.dump(out, sort_keys=False, default_flow_style=False)
    dst = CONFIGS / dst_dir / dst_name
    dst.parent.mkdir(exist_ok=True)
    (dst.parent / "__init__.py").touch()
    dst.write_text(header + body + SIM_BEAMLINE_DEVICES)
    print(f"{dst_name}: {len(out)} devices, dropped: {dropped}")


if __name__ == "__main__":
    transform("ptycho_lamni.yaml", "simulated_omny", "simulated_lamni.yaml", "LamNI")
    transform("ptycho_omny.yaml", "simulated_omny", "simulated_omny.yaml", "OMNY")