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 initLog():
    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)
    return log

def feedbackLoop(apparatus, nsteps):
    # while True:
    for i in range(0,nsteps):
        adjustRFandLog(apparatus)
    apparatus.runStatus = True

def adjustRFandLog(apparatus):
    timenow = time.time()

    scope = apparatus.scope
    rfgen = apparatus.rfgen
    pid   = apparatus.pid
    log   = apparatus.log

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

def updateLogPlot(apparatus):
    log   = apparatus.log
    plt.clf()
    log.plot()
    plt.pause(.0001)
    print(f're-plotting')

def updateLogPlotLoop(apparatus):
    while True:
        updateLogPlot(apparatus)
        if not apparatus.runStatus:
            break

def updateLogFile(apparatus):
    log   = apparatus.log
    logFileName = apparatus.logFileName
    if logFileName is not None:
        log.save(logFileName, item_format='.15e')
        print(f'Saving to {logFileName}')

def updateLogFileLoop(apparatus):
    while True:
        updateLogFile(apparatus)
        if not apparatus.runStatus:
            break


freqZero = 6834687190
freqDeltaM2 = 6835387164
def 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()
    pid = PID(100,400,0, sign=-1)

    apparatus.scope = scope
    apparatus.rfgen = rfgen
    apparatus.pid   = pid
    apparatus.log   = initLog()
    apparatus.runStatus = False



    SweepSpan = 10000
    apparatus.rfgen.stopFrequencySweep()
    apparatus.rfgen.setFreqFixed(freqDeltaM2+100)
    apparatus.logFileName = apparatus.getNextDataFile()
    print(apparatus.logFileName)


    # initial lock
    apparatus.logFileName = None
    apparatus.pid.reset()
    apparatus.runStatus = True
    feedbackLoop(apparatus, nsteps=10)
    updateLogPlot(apparatus)

    # long term lock
    apparatus.logFileName = None
    apparatus.pid.reset()
    apparatus.runStatus = True
    feedbackLoop(apparatus, nsteps=100)
    updateLogPlot(apparatus)

    # apparatus.log.save(apparatus.logFileName, item_format='.15e')
    print(apparatus.logFileName)
    return(apparatus)



if __name__ == '__main__': 
    mpl.use("TkAgg")
    apparatus=main()