import justpy as jp
import pyvisa
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

import asyncio
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

async def feedbackLoop(apparatus, nsteps):
    # while True:
    for i in range(0,nsteps):
        adjustRFandLog(apparatus)
        await asyncio.sleep(0.001)
    apparatus.runStatus = False

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
    f = plt.figure()
    log.plot()
    log_chart.set_figure(f)
    plt.close(f)
    print(f're-plotting')
    st.text=f're-plotting {time.time()}'
    jp.run_task(wp.update())

async def updateLogPlotLoop(apparatus):
    while True:
        updateLogPlot(apparatus)
        if not apparatus.runStatus:
            break
        await asyncio.sleep(5)

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


async def updateLogFileLoop(apparatus):
    while True:
        updateLogFile(apparatus)
        if not apparatus.runStatus:
            break
        await asyncio.sleep(10)


freqZero = 6834687190
freqDeltaM2 = 6835387164

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


    print('========== Initial lock ===========')
    apparatus.logFileName = None
    apparatus.pid.reset()
    apparatus.runStatus = True
    res = await asyncio.gather( feedbackLoop(apparatus, nsteps=100), updateLogPlotLoop(apparatus), updateLogFileLoop(apparatus) )
    print(apparatus.logFileName)
    updateLogPlot(apparatus)

    print('========= Long term lock ===========')
    apparatus.log   = initLog()
    apparatus.logFileName = None
    apparatus.logFileName = apparatus.getNextDataFile()
    apparatus.pid.reset()
    apparatus.runStatus = True
    res = await asyncio.gather( feedbackLoop(apparatus, nsteps=1000), updateLogPlotLoop(apparatus), updateLogFileLoop(apparatus) )
    print(apparatus.logFileName)
    updateLogPlot(apparatus)

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


async def getPage():
    return wp

async def  jp_startup():
    jp.run_task(main())

if __name__ == '__main__': 
    wp = jp.WebPage(delete_flag=False)
    d=jp.Div(text='Magnetometer lock', a=wp)
    f = plt.figure()
    log_chart = jp.Matplotlib(a=wp)
    plt.close(f)
    # mpl.use("TkAgg")
    # apparatus = asyncio.run(main())
    st = jp.Div(test='Status string', a=wp)
    jp.justpy(getPage, startup=jp_startup)