import pyvisa
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import time

from qolab.hardware import BasicInstrument
from qolab.hardware.scope import SDS1104X
from qolab.hardware.rf_generator import AgilentE8257D
from qolab.feedback import PID
from qolab.data import TraceSetSameX, TraceXY, Trace

def adjustRF(pid, scope, rfgen, nsteps=1000, fname=None):
    pid.reset()

    errorTrace = Trace("error")
    errorTrace.config['unit']='V'
    timeTrace = Trace("time")
    timeTrace.config['unit']='S'
    errorLog = TraceXY("error")
    errorLog.x = timeTrace
    errorLog.y = errorTrace

    freqTrace = Trace("frequency")
    freqTrace.config['unit']='Hz'
    freqLog = TraceXY("frequency")
    freqLog.x = timeTrace
    freqLog.y = freqTrace

    feedbackTrace = Trace("feedback")
    feedbackTrace.config['unit']='Hz'
    feedbackLog = TraceXY("feedback")
    feedbackLog.x = timeTrace
    feedbackLog.y = feedbackTrace


    log = TraceSetSameX("timelog")
    log.addTrace(errorLog)
    log.addTrace(freqLog)
    log.addTrace(feedbackLog)


    # while True:
    for i in range(0,nsteps):
        timenow = time.time()

        ch1=scope.getTrace(1, decimate=False)
        err = np.mean(ch1.y.values)
        fdbck = pid.feedback(err)

        freq0 = rfgen.getFreqFixed()
        freq = freq0 + fdbck
        rfgen.setFreqFixed(freq)
        print(f'error = {err}, feedback = {fdbck}, request freq = {freq}')

        log.addPointToTrace(timenow)
        log.addPointToTrace(err, "error")
        log.addPointToTrace(freq0, "frequency")
        log.addPointToTrace(fdbck, "feedback")
        
        if fname is not None:
            log.save(fname, item_format='.15e')

    plt.clf()
    log.plot()
    return log



if __name__ == '__main__': 
    apparatus = BasicInstrument()
    apparatus.config['Device type'] = 'QOL VAMPIRE HighPower magnetometer'
    apparatus.config['Device model'] = 'v0.1'
    apparatus.config['FnamePrefix'] = 'locked_magnetometer'
    apparatus.config['SavePath'] = '/mnt/qol_grp_data/data.VAMPIRE.HighPower'


    print("testing")
    rm = pyvisa.ResourceManager()
    instr_scope=rm.open_resource('TCPIP::192.168.0.61::INSTR')
    scope = SDS1104X(instr_scope)
    instr_rfgen=rm.open_resource('TCPIP::192.168.0.114::INSTR')
    rfgen = AgilentE8257D(instr_rfgen)
    print('------ Header start -------------')
    print(str.join('\n', scope.getHeader()))
    print(str.join('\n', rfgen.getHeader()))
    print('------ Header ends  -------------')
    ch1  = scope.getTrace(1)
    traces = scope.getAllTraces()

    freqZero = 6834687190
    freqDeltaM2 = 6835387164
    SweepSpan = 10000

    rfgen.stopFrequencySweep()
    rfgen.setFreqFixed(freqDeltaM2+100)

    pid = PID(100,400,0, sign=-1)

    # initial lock
    log=adjustRF(pid, scope, rfgen, nsteps=100)

    # long term lock
    fname = apparatus.getNextDataFile()
    print(fname)
    log=adjustRF(pid, scope, rfgen, nsteps=10000, fname=fname)

    log.save(fname, item_format='.15e')