""" DilSc prototype

"""

from slic.core.adjustable import Adjustable, PVAdjustable

from slic.core.device import Device, SimpleDevice

from frappy.client import SecopClient
from frappy import states
from frappy.datatypes import StatusType


class Dilution(Device):
    def __init__(self, **kwargs):

        self.name = 'DilSc'
        ID = self.name
        super().__init__(ID, **kwargs)

        self.address = 'dilsc.psi.ch:5000'
        self.dilsc = SecopClient(self.address)
        self.dilsc.connect()
        
        self.x = MagnetCoil("X", self.dilsc, 'x', limit_low=-0.6, limit_high=0.6)
        self.y = MagnetCoil("Y", self.dilsc, 'y', limit_low=-0.6, limit_high=0.6)
        self.z = MagnetCoil("Z", self.dilsc, 'z', limit_low=-5.2, limit_high=5.2)
        self.T_plato = Thermometer('T_plato', self.dilsc, limit_low=0, limit_high=300) 
        self.T_chip = Thermometer('T_chip', self.dilsc, limit_low=0, limit_high=300)
        self.T_reg = Thermometer('T_reg', self.dilsc, limit_low=0, limit_high=300)

class Thermometer(Adjustable):

    def __init__(self, name, dilsc_connection, limit_low=-0.0001, limit_high=0.0001):
        
        super().__init__(name, limit_low=limit_low, limit_high=limit_high)
        self.dilsc = dilsc_connection

        # Heater to regulate channel A on the Galgen LakeShore372. Give value in Ohm.
        self.heater_resistance = 82
        # Defines heater ranges and PID values for various control temperature regions.  
        self.heater_settings = [
            {'range': 'off', 'pid': (800, 1, 0), 'temp_start': -1, 'temp_end': 0.00099},
            {'range': '300uA', 'pid': (800, 1, 0), 'temp_start': 0.001, 'temp_end': 0.037},
            {'range': '1mA', 'pid': (40, 14, 0), 'temp_start': 0.038, 'temp_end': 0.108},
            {'range': '3mA', 'pid': (4, 15, 0), 'temp_start': 0.110, 'temp_end': 0.350},
            {'range': '10mA', 'pid': (2, 12, 0), 'temp_start': 0.352, 'temp_end': 1.100},
            {'range': '30mA', 'pid': (2, 60, 0), 'temp_start': 1.100, 'temp_end': 2.300}
            ]
            
    def _check_connection(func):
        def checker(self, *args, **kwargs):
            if not self.dilsc.online:
                raise ConnectionError(f'No connection to dilsc at {self.address}')
            else:
                return func(self, *args, **kwargs)
        return checker

    @_check_connection
    def set_heater_range(self, range_value): 
        """ Sets the heater range for the control channel and sets the PID parameters for the associated control loop. 
            TODO: Consider redoing so that it only works for the T_reg channel.
            TODO: the pid values are now sent with stringio and not saved as parameters in slic, this should be fixed.
            - 
        """
        heater_range_map = {
            'off': 0, '30uA': 1, '100uA': 2, '300uA': 3, '1mA': 4, '3mA': 5, '10mA': 6, '30mA': 7, '100mA': 8}
        if range_value not in heater_range_map:
            raise ValueError("Invalid range value. Allowed are: ['off','30uA','100uA','300uA','1mA','3mA','10mA','30mA','100mA']")
        enum_value = heater_range_map[range_value]
        self.dilsc.setParameter(self.name, 'htrrng', enum_value)
        
        for setting in self.heater_settings:
            if setting['range'] == range_value:
                self.set_PID_parameters(*setting['pid'])
                break
        return 

    @_check_connection
    def get_current_value(self):
        cacheitem = self.dilsc.getParameter(f'{self.name}', 'value', trycache=False)
        return cacheitem.value
        
    
    @_check_connection
    def set_target_value(self, value, adjust_heater_range=True):
        if adjust_heater_range:
            range_to_set = None
            for setting in self.heater_settings:
                if setting['temp_start'] <= value and setting['temp_end'] >= value:
                    range_to_set = setting['range']
            if range_to_set != None:
                self.set_heater_range(range_to_set)
            else: 
                raise ValueError(f"Heater range for the target value {value} could not be set")
        self.dilsc.setParameter(f'{self.name}', 'target', value)

    @_check_connection
    def get_target_value(self):
        cacheitem = self.dilsc.getParameter(f'{self.name}', 'target', trycache=False)
        return cacheitem.value

            
    @_check_connection
    def is_moving(self):
        response = self.dilsc.getParameter(f'{self.name}','status', trycache=False)
        return response[0][0] > StatusType.PREPARED
    
    @_check_connection
    def get_PID_parameters(self):
        """ Returns the current PID parameters associated with the control loop for 
            this thermometer.
        """
        # response = self.dilsc.getParameter(f'{self.name}','ctrlpars', trycache=False)
        response = self.dilsc.execCommand('lscio', 'communicate', 'PID?')
        return response
    
    @_check_connection
    def set_PID_parameters(self, p, i, d):
        """ Sets the PID parameters for the associated control loop. 
            TODO:
            - This still returns a timeout error but sets the correct values.
            - The range is limited to less than the Lakeshore range allows, this needs
              to be fixed in frappy.
        """
        # self.dilsc.setParameter(f'{self.name}', 'ctrlpars', {'p': p, 'i': i, 'd': d})
        command_string = f'PID {0},{p},{i},{d}; PID?'
        self.dilsc.execCommand('lscio', 'communicate', command_string)
        

class MagnetCoil(Adjustable):

    def __init__(self, name, dilsc_connection, direction, limit_low=-0.0001, limit_high=0.0001):
        
        super().__init__(name, limit_low=-0.0001, limit_high=0.0001)
        
        self.direction = direction.lower()    

        if self.direction not in ["x", "y", "z"]:
            raise ValueError("Direction must be either x, y or z.")
        
        self.dilsc = dilsc_connection
        
            
    def _check_connection(func):
        def checker(self, *args, **kwargs):
            if not self.dilsc.online:
                raise ConnectionError(f'No connection to dilsc at {self.address}')
            else:
                return func(self, *args, **kwargs)
        return checker

    @_check_connection
    def get_current_value(self):
        cacheitem = self.dilsc.getParameter(f'mf{self.direction}', 'value', trycache=False)
        return cacheitem.value
        
    
    @_check_connection
    def set_target_value(self, value):
        self.dilsc.setParameter(f'mf{self.direction}', 'target', value)
            
    @_check_connection
    def is_moving(self):
        response = self.dilsc.getParameter(f'mf{self.direction}','status', trycache=False)
        return response[0][0] > StatusType.PREPARED
    
    @_check_connection
    def set_ramp_speed(self, direction, value):
        """ Sets ramp speed for a given direction (x,y or z) and value in T/min.
        """
        if value > 0.5:
            raise ValueError('Do not exceed 0.5 T/min unless you want a quench party')
        if value <= 0:
            raise ValueError('Only positive values are allowed')
        if self.direction not in ["x", "y", "z"]:
            raise ValueError("Direction must be either x, y or z.")
        else:
            self.dilsc.setParameter(direction, 'ramp', value)
        return

    