import pyvisa
import time

class KeysightU2042XA:
    
    def __init__(self,adress,freq,cal=1):
        #establish connection to device with pyvisa. The device is initiallized with visa adress "adress" and with frequency "freq". If cal == 1 a on e time calibration
        #and zeroing is performed
        self.rm = pyvisa.ResourceManager()
        try:
            self.inst = self.rm.open_resource(adress)
        except:
            print(f'could not connect to {adress}')
            return
        self.inst.timeout = 20000
        self.inst.write('SYST:PRES')
        # time.sleep(0.2)
        self.inst.write(f'SENS:FREQ {freq}')
        time.sleep(0.2)
        if cal == 1:
            self.inst.write('CAL:ZERO:AUTO ONCE')
            time.sleep(16)
            self.inst.write('CAL:AUTO ONCE')
            time.sleep(6)
        self.inst.write('CAL:ZERO:AUTO OFF')
        self.inst.write('CAL:AUTO OFF')
        self.inst.write('INIT:CONT ON')
        time.sleep(2)
        
    def __del__(self):
        self.rm.close()

    def activate_CW(self):
        #bring the device in a state to measure CW. The values can then be obtained by "read"
        self.inst.write('CAL:ZERO:AUTO OFF')
        self.inst.write('CAL:AUTO OFF')
        self.inst.write('DET:FUNC AVER')
        self.inst.write('INIT:CONT ON')

    def activate_trace(self):
        #brings the device in a state to measure traces. The traces are obtained with read_trace
        self.inst.write('CAL:ZERO:AUTO OFF')
        self.inst.write('CAL:AUTO OFF')
        self.inst.write('DET:FUNC NORM')
        self.inst.write('INIT:CONT OFF')
        self.inst.write('TRIG:SOUR INT')
        self.inst.write('TRAC:STAT ON')    

    def read(self):
        #reads the CW or Pulse measurement from the device
        data = float(self.inst.query('FETC?'))
        return data
    
    def read_trace(self,res):
        #reads a trace with resolution "res" from the device
        self.inst.write('INIT')
        # data = self.inst.query(f"TRAC? {res}")
        data = self.inst.query_binary_values(f"TRAC? {res}", datatype='f', is_big_endian = True)
        # self.inst.write('TRAC? LRES')
        # data = self.inst.read_raw()
        return data
    
    def set_trace_time(self,t):
        #sets the duration of the trace
        self.inst.write(f"SENS:TRAC:TIME {t}")

    def set_trigger_delay(self,t):
        #sets the trigger delay of the trace
        print(f'TRIG:DEL {t}')
        self.inst.write(f'TRIG:DEL {t}')

    def set_trigger_level(self,*args):
        #sets the trigger level of the device. If no argument is given it is set to auto
        if args:
            self.inst.write(f"TRIG:LEV:AUTO 0")
            self.inst.write(f"TRIG:LEV {args[0]}")
        else:
            self.inst.write(f"TRIG:LEV:AUTO 1")