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'

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

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

    def setFreq(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)

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(f'Fixed frequency: {rfgen.getFreq()}')
    print(f'Sweep Central frequency: {rfgen.getSweepCentralFreq()}')
    print(f'Sweep Span: {rfgen.getSweepSpan()}')