from qolab.hardware.basic import BasicInstrument
from qolab.hardware.power_supply import PowerSupplySCPI
import time

class KeysightE3612A(PowerSupplySCPI):
    """ Keysight E3612A power supply """
    def __init__(self, resource, *args, **kwds):
        super().__init__(resource, *args, **kwds)
        self.resource.read_termination='\n'
        self.config['Device model'] = 'Keysight E3612A'
        self.numberOfChannels = 3
        self.deviceProperties = {'OpMode', };
        self.channelProperties = {'IsOn', 'Regulation', 'Vout', 'Vlimit', 'Iout', 'Ilimit', }

    def getOpMode(self):
        """
        Queries power supply operation mode, returns OFF|PAR|SER|TRAC
        OFF stands for independent channels
        """
        qstr = f'OUTP:PAIR?'
        rstr = self.query(qstr)
        return( rstr )

    def setOpMode(self, val):
        """
        Sets power supply operation mode, returns OFF|PAR|SER|TRAC
        OFF stands for independent channels
        """
        cmnd = f'OUTP:PAIR {val}'
        rstr = self.write(cmnd)

    def setChanOn(self, chNum):
        """ Power up channel output """
        self.write(f'OUTP ON,(@{chNum})')

    def setChanOff(self, chNum):
        """ Power down channel output """
        self.write(f'OUTP OFF,(@{chNum})')

    @BasicInstrument.tsdb_append
    def getChanIsOn(self, chNum):
        """ Queries channel output state """
        qstr = f'OUTP? (@{chNum})'
        rstr = self.query(qstr)
        return( bool(float(rstr)) )

    @BasicInstrument.tsdb_append
    def getChanRegulation(self, chNum):
        """
        Queries channel output regulation
        0 - The output is off and unregulated
        1 - The output is CC (constant current) operating mode
        2 - The output is CV (constant voltage) operating mode
        3 - The output has hardware failure
        """
        qstr = f'STAT:QUES:INST:ISUM{chNum}:COND?'
        rstr = self.query(qstr)
        return( int(rstr) )

    @BasicInstrument.tsdb_append
    def getChanVout(self, chNum):
        qstr = f'MEAS:VOLT? (@{chNum})'
        rstr = self.query(qstr)
        return( float(rstr) )

    @BasicInstrument.tsdb_append
    def getChanVlimit(self, chNum):
        qstr = f'SOUR:VOLT? (@{chNum})'
        rstr = self.query(qstr)
        return( float(rstr) )

    @BasicInstrument.tsdb_append
    def setChanVlimit(self, chNum, val):
        if chNum == 1 and val > 6.180:
            val = 6.180
        if (chNum == 2 or chNum == 3) and val > 25.750:
            val = 25.750
        cmnd = f'SOURCe:VOLT {val},(@{chNum})'
        rstr = self.write(cmnd)

    @BasicInstrument.tsdb_append
    def getChanIout(self, chNum):
        qstr = f'MEAS:CURR? (@{chNum})'
        rstr = self.query(qstr)
        return( float(rstr) )

    @BasicInstrument.tsdb_append
    def setChanIout(self, chNum, val, currentPrecision=5e-6, currentHeadRoom=1e-3, dwellTime=0.3):
        """
        Tuning Vout to achieve desired Iout.
        Generally setting current limit will maintain current near but not exact to desired.
        This function will try to guess the correct Vout value
        """
        iDesired = val
        self.setChanIlimit(chNum, val+currentHeadRoom)
        # here we assume that hook up is already made, so we can estimate source resistance
        for i in range(10):
            iOut=self.getChanIout(chNum)
            if abs(iOut-iDesired) <= currentPrecision:
                break
            vOut=self.getChanVout(chNum)
            R=vOut/iOut
            vDesired = R*iDesired
            self.setChanVlimit(chNum, vDesired)
            time.sleep(dwellTime)

    @BasicInstrument.tsdb_append
    def getChanIlimit(self, chNum):
        qstr = f'SOURce:CURR? (@{chNum})'
        rstr = self.query(qstr)
        return( float(rstr) )

    @BasicInstrument.tsdb_append
    def setChanIlimit(self, chNum, val):
        """ Set current limit, seems to be >=0.002 for Ch1 and >=0.001 for Ch2 and Ch3 """
        if chNum == 1 and val < 0.002:
            val = 0.002
        if chNum == 1 and val > 5.150:
            val = 5.150
        if (chNum == 2 or chNum == 3) and val < 0.001:
            val = 0.001
        if (chNum == 2 or chNum == 3) and val > 1.030:
            val = 1.030
        cmnd = f'SOURCe:CURR {val},(@{chNum})'
        rstr = self.write(cmnd)


if __name__ == '__main__': 
    import pyvisa
    print("testing")
    rm = pyvisa.ResourceManager()
    print(rm.list_resources())    
    instr=rm.open_resource('USB0::10893::4354::MY61001869::0::INSTR')
    ps = KeysightE3612A(instr)
    print('------ Header start -------------')
    print(str.join('\n', ps.getHeader()))
    print('------ Header ends  -------------')