from aare.common.coordinate import Coordinate
from aare.common.models import StagePositionEnum
from aare.daq.devices import BeamlineDevices
from aare.daq.config import BeamlineConfig, ABR_POS_MOUNT


def move_bsz(devs: BeamlineDevices, target: float):
    if abs(target - devs.bsz.position) > 0.1:
        beamstop_stage_measure = devs.beamstop_stage_up
        reflector_measure = devs.reflector_up
        devs.reflector_up = False
        devs.beamstop_stage_up = True
        devs.beamstop_z = target

        if reflector_measure:
            devs.reflector_up = True

        if not beamstop_stage_measure:
            devs.beamstop_stage_up = False


def wait_for_dc_devices(devs: BeamlineDevices):
    # Should wait for DTZ
    pass

def wait_for_se_devices(devs: BeamlineDevices, cfg: BeamlineConfig):
    # Should wait for Cryojet go far
    pass

def common_2rse(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 = Coordinate(x=0.0, y=0.0, z=0.0)
    devs.aerotech_omega = 0.0
    #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

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 = Coordinate(x=0.0, y=0.0, z=0.0)
    devs.aerotech_omega = 0.0
    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

def sa2rse(devs: BeamlineDevices, cfg: BeamlineConfig):
    common_2rse(devs, cfg)

def dc2rse(devs: BeamlineDevices, cfg: BeamlineConfig):
    common_2rse(devs, cfg)

def se2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    devs.reflector_up = StagePositionEnum.MEASURE
    devs.aerotech_pos = cfg.abr_meas_pos

def rse2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    se2sa(devs, cfg)

def sa2dc(devs: BeamlineDevices, cfg: BeamlineConfig):
    pass

def dc2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    pass

def sa2xrf(devs: BeamlineDevices, cfg: BeamlineConfig):
    """sample alignment to XrfCollection"""
    pass


def xrf2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    """XrfCollection to sample alignment"""
    pass


def sa2ws(devs: BeamlineDevices, cfg: BeamlineConfig):
    pass


def ws2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    pass


def sa2ba(devs: BeamlineDevices, cfg: BeamlineConfig):
    pass


def ba2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    pass

def sa2bl(devs: BeamlineDevices, cfg: BeamlineConfig):
    devs.scintillator = 20.0

def bl2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    devs.scintillator = 20.0


def bl2ba(devs: BeamlineDevices, cfg: BeamlineConfig):
    devs.scintillator = 20.0

def ba2bl(devs: BeamlineDevices, cfg: BeamlineConfig):
    devs.scintillator = 20.0

def sa2dh(devs: BeamlineDevices, cfg: BeamlineConfig):
    devs.scintillator = 20.0
    if devs.tell.get_mounted_sample() is not None:
        try:
            # Best effort try to unmount
            devs.tell.unmount(wait=True)
        except Exception as e:
            print(f"Error for unmounting: {e}")
    devs.tell.dry(wait_cold=-1, wait=False)

def dh2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    devs.scintillator = 20.0
    devs.reflector_up = StagePositionEnum.MEASURE

if __name__ == "__main__":
    from aare.common.beamline import mx_beamline
    import time
    bl = mx_beamline()
    cfg = BeamlineConfig(bl)
    devs = BeamlineDevices(bl)
    common_2rse(devs, cfg)
    time.sleep(10.0)
    rse2sa(devs, cfg)