from qolab.hardware.rf_generator import RFGeneratorSCPI
from qolab.data.trace import Trace
import re
import numpy as np

class AgilentE8257D(RFGeneratorSCPI):
    """ Agilent E8257D RF generator """
    def __init__(self, resource):
        super().__init__(resource)
        self.resource.read_termination='\n'
        self.config['Device model'] = 'Agilent E8257D'
        self.deviceProperties.extend([
            'RFPowerState',
            'RFAmplitude',
            'ModulationState',
            'SweepCentralFreq', 'SweepSpan',
            ])

    """ Note: Fixed frequency and Center frequency (of sweep)  are different in this model """

    def getFreqFixed(self):
        qstr = f':FREQuency:Fixed?'
        rstr = self.query(qstr)
        return( float(rstr) )

    def setFreqFixed(self, freq):
        cstr = f':FREQuency:FIXED {freq}Hz'
        self.write(cstr)

    def getSweepCentralFreq(self):
        qstr = f':FREQuency:CENTer?'
        rstr = self.query(qstr)
        return( float(rstr) )

    def setSweepCentralFreq(self, cfreq):
        cstr = f':FREQuency:CENTer {cfreq}Hz'
        self.write(cstr)

    def getSweepSpan(self):
        qstr = f':FREQuency:SPAN?'
        rstr = self.query(qstr)
        return( float(rstr) )

    def setSweepSpan(self, span):
        cstr = f':FREQuency:SPAN {span}Hz'
        self.write(cstr)

    def setSweep(self, cfreq=None, span=None):
        if cfreq is not None:
            self.setSweepCentralFreq(cfreq)
        if cfreq is not None:
            self.setSweepSpan(span)

    def getModulationState(self):
        return(int(self.query(f':MODulation:STATe?')))

    def setModulationState(self, val):
        rstr = self.write(f':MODulation:STATe {val}')

    def getRFPowerState(self):
        return(int(self.query(f':OUTPut:STATe?')))

    def setRFPowerState(self, val):
        rstr = self.write(f':OUTPut:STATe {val}')

    def getRFAmplitude(self):
        return(float(self.query(f':POWer:AMPLitude?')))

    def seRFAmplitude(self, val):
        rstr = self.write(f':POWer:AMPLitude {val}')

if __name__ == '__main__': 
    import pyvisa
    print("testing")
    rm = pyvisa.ResourceManager()
    print(rm.list_resources())    
    instr=rm.open_resource('TCPIP::192.168.0.114::INSTR')
    rfgen = AgilentE8257D(instr)
    print('------ Header start -------------')
    print(str.join('\n', rfgen.getHeader()))
    print('------ Header ends  -------------')