from qolab.hardware.basic import BasicInstrument
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, *args, **kwds):
        super().__init__(resource, *args, **kwds)
        self.resource.read_termination='\n'
        self.config['Device model'] = 'Agilent E8257D'
        self.deviceProperties.extend([
            'RFPowerState',
            'RFAmplitude',
            'ModulationState',
            'FM1State',
            'FM1Source',
            'FM1ModulationDepth',
            'FM2State',
            'FM2Source',
            'FM2ModulationDepth',
            'FrequencyMode', # sweep or continious
            'SweepCentralFreq', 'SweepSpan',
            ])

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

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

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

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

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

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

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

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

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

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

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

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

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

    @BasicInstrument.tsdb_append
    def setRFAmplitude(self, val):
        rstr = self.write(f':POWer:AMPLitude {val}')

    @BasicInstrument.tsdb_append
    def getFM1ModulationDepth(self):
        return(float(self.query(f':FM1:Deviation?')))

    @BasicInstrument.tsdb_append
    def setFM1ModulationDepth(self, val):
        rstr = self.write(f':FM1:Deviation {val}')

    @BasicInstrument.tsdb_append
    def getFM2ModulationDepth(self):
        return(float(self.query(f':FM2:Deviation?')))

    @BasicInstrument.tsdb_append
    def setFM2ModulationDepth(self, val):
        rstr = self.write(f':FM2:Deviation {val}')

    @BasicInstrument.tsdb_append
    def getFM1Source(self):
        return(str(self.query(f':FM1:Source?')))

    @BasicInstrument.tsdb_append
    def setFM1Source(self, val):
        rstr = self.write(f':FM1:Source {val}')

    @BasicInstrument.tsdb_append
    def getFM2Source(self):
        return(str(self.query(f':FM2:Source?')))

    @BasicInstrument.tsdb_append
    def setFM2Source(self, val):
        rstr = self.write(f':FM2:Source {val}')

    @BasicInstrument.tsdb_append
    def getFM1State(self):
        return(int(self.query(f':FM1:State?')))

    @BasicInstrument.tsdb_append
    def setFM1State(self, val):
        rstr = self.write(f':FM1:State {val}')

    @BasicInstrument.tsdb_append
    def getFM2State(self):
        return(int(self.query(f':FM2:State?')))

    @BasicInstrument.tsdb_append
    def setFM2State(self, val):
        rstr = self.write(f':FM2:State {val}')

    @BasicInstrument.tsdb_append
    def getFrequencyMode(self):
        return(str(self.query(f':Frequency:Mode?')))

    @BasicInstrument.tsdb_append
    def setFrequencyMode(self, val):
        rstr = self.write(f':Frequency:Mode {val}')

    @BasicInstrument.tsdb_append
    def startFrequencySweep(self):
        self.setFrequencyMode('Sweep')

    def stopFrequencySweep(self):
        self.setFrequencyMode('FIXED')

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  -------------')