""" Diffractometer

    motorized axis: 

    SARES30-CPCL-ECMC02:ROT2THETA
    SARES30-CPCL-ECMC02:ROTTHETA
    SARES30-CPCL-ECMC02:TRXBASE
    SARES30-CPCL-ECMC02:TRYBASE

"""

import time

from slic.core.adjustable import PVAdjustable, Linked, Adjustable, Collection
from slic.core.device import Device
from slic.devices.general.motor import Motor

class MotorECMC(Motor):
    '''Regular motor class with extras for ECMC servo motors (tolerance deadbands, PIDs ...)'''
    def __init__(self, ID, **kwargs):
        super().__init__(ID, **kwargs)
        self.deadband_epics = PVAdjustable(self.name+'-SetAtTrgtDB',units='mm',name='Tolerance deadband for EPICS target reached flag.')
        self.deadband_ecmc  = PVAdjustable(self.name+'-SetCtrlDB',units='mm', name='Tolerance deadband for ECMC servo motor control')
        self.deadbands      = Collection(self.name+'motion deadbands',[self.deadband_epics,self.deadband_ecmc])
        
        self.p = PVAdjustable(self.name+'-Ctrl-Kp',name='P parameter')
        self.i = PVAdjustable(self.name+'-Ctrl-Ki',name='I parameter')
        self.d = PVAdjustable(self.name+'-Ctrl-Kd',name='D parameter')
        self.pids  = Collection(self.name+'PID values',[self.p,self.i,self.d])
        self.enabled = PVAdjustable(self.name+'-EnaCmd',name='Axis enabled')

class Diffractometer(Device):
    def __init__(self, ID, **kwargs):
        super().__init__(ID, **kwargs)
        
        self.twotheta = MotorECMC(ID + ":ROT2THETA")  # , ID=None, name=None, units=None, internal=False):
        self.theta = MotorECMC(ID + ":ROTTHETA")      # , ID=None, name=None, units=None, internal=False):
        
        self.trx_base = MotorECMC(ID + ":TRXBASE")     # , ID=None, name=None, units=None, internal=False):
        self.try_base = MotorECMC(ID + ":TRYBASE")     # , ID=None, name=None, units=None, internal=False):

        self.tr_x = MotorECMC(ID + ":TRX")
        self.tr_y = MotorECMC(ID + ":TRY")
        self.tr_z = MotorECMC(ID + ":TRZ")

        self.td = MotorECMC(ID + ":TD")
        self.td_y = Motor("SARES30-MOBI2:MOT_Z")   # Extra Huber stage for the diffractometer motor

        if ID == "SARES32-GPS":
            self.name = "DM2: Cristallina pulsed magnet diffractometer"
            # This diffractometer also has extra swivel stages
            self.rot_x = MotorECMC(ID + ":ROTX")
            self.rot_z = MotorECMC(ID + ":ROTZ")
            self.axes = Collection(self.name+'movable axes',[self.twotheta,self.theta,self.tr_x,self.tr_y,self.tr_z,self.td,self.trx_base,self.try_base,
                                                           self.rot_x,self.rot_z])
        if ID == "SARES31-GPS":
            self.name = "DM1: Cristallina dilution fridge diffractometer"
            self.axes = Collection(self.name+'movable axes',[self.twotheta,self.theta,self.tr_x,self.tr_y,self.tr_z,self.td,self.trx_base,self.try_base])

        
    def enable_all_axes(self):
        for axis in self.axes.adjs:
            axis.enabled.set_target_value(True).wait()

    def disable_all_axes(self):
        for axis in self.axes.adjs:
            axis.enabled.set_target_value(False).wait()
        

# class Diffractometer(Device):
#     def __init__(self, ID, motor_naming="MOTOR", **kwargs):
#         super().__init__(ID, **kwargs)
        
#         self.twotheta = Motor("SARES30-CPCL-ECMC02:ROT2THETA")  # , ID=None, name=None, units=None, internal=False):
#         self.theta = Motor("SARES30-CPCL-ECMC02:ROTTHETA")      # , ID=None, name=None, units=None, internal=False):
        
#         self.trx_base = Motor("SARES30-CPCL-ECMC02:TRXBASE")     # , ID=None, name=None, units=None, internal=False):
#         self.try_base = Motor("SARES30-CPCL-ECMC02:TRYBASE")     # , ID=None, name=None, units=None, internal=False):

#         self.tr_x = Motor("SARES30-CPCL-ECMC02:TRX")
#         self.tr_y = Motor("SARES30-CPCL-ECMC02:TRY")
#         self.tr_z = Motor("SARES30-CPCL-ECMC02:TRZ")

#         self.td = Motor("SARES30-CPCL-ECMC02:TD")
        

# Set speed: 
# diffractometer.theta._motor.VELO = 0.25

class ThetasCombined(Linked):
    def __init__(self, *args, **kwargs):
        super().__init__(self, *args, **kwargs)

    def connect_axis(self):
        """
         calculate offset to match scale factor
        """

        offset = self.secondary.get_current_value() - self.primary.get_current_value() * self.scale_factor
        self.offset = offset


class DelayAxis(MotorECMC):
    def __init__(self, axis, delay_after_move=1):
        #super().__init__(self, *args, **kwargs)#
        delay_time = delay_after_move
        self.axis = axis

    def set_target_value(self, value):
        start_time = time.time()

        self.axis.set_target_value(value)
        while self.axis.is_moving():
            time.sleep(0.1)

        elapsed = time.time() - start_time
        wait_time = max(0, self.delay_time - elapsed)
        time.sleep(wait_time)


class Delayed(Adjustable):
    def __init__(self, ID, adj, delay_after_move=0, **kwargs):
        super().__init__(ID, **kwargs)
        self.adj = adj
        self.delay_after_move = delay_after_move

    def get_current_value(self):
        value = self.adj.get_current_value()
        return value
    
    def set_target_value(self, value):
        start_time = time.time()

        self.adj.set_target_value(value).wait()

        while self.adj.is_moving():
            time.sleep(0.02)

        elapsed = time.time() - start_time
        wait_time = max(0, self.delay_time - elapsed)
        time.sleep(wait_time)

    def is_moving(self):
        return self.adj.is_moving()