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

import asyncio
import time

import qolab.gui.web as gui
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.1)
    apparatus.runStatus = False

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

    scope = apparatus.scope
    rfgen = apparatus.rfgen
    pid   = apparatus.pid
    log   = apparatus.gui_log.traces

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


freqZero = 6834686400
freqDeltaM2 = 6835387100

async def main():
    task_wp_update_loop = asyncio.create_task(update_webpage_loop(update_interval=1))
    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'
    apparatus.config['SavePath'] = './data'

    apparatus.gui_log = gui.QOLTimeLog(a=wp)
    apparatus.gui_log.save_controls.getNextDataFile = apparatus.getNextDataFile

    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.state = None
    apparatus.runStatus = False



    SweepSpan = 10000
    apparatus.rfgen.stopFrequencySweep()
    apparatus.rfgen.setFreqFixed(freqDeltaM2)
    # apparatus.rfgen.setFreqFixed(freqZero)


    print('========== Initial lock ===========')
    apparatus.state = 'Initial lock'
    update_webpage(apparatus=apparatus)
    apparatus.gui_log.setTraces( initLog() )
    apparatus.pid.reset()
    apparatus.runStatus = True
    res = await asyncio.gather(
            feedbackLoop(apparatus, nsteps=100)
            )

    print('========= Long term lock ===========')
    apparatus.state = 'Long term lock'
    update_webpage(apparatus=apparatus)
    apparatus.gui_log.setTraces( initLog() )
    apparatus.pid.reset()
    apparatus.runStatus = True
    res = await asyncio.gather(
            feedbackLoop(apparatus, nsteps=100000)
            )

    apparatus.state = 'Done working with hardware'

    apparatus.gui_log.stop_tasks()
    task_wp_update_loop.cancel()
    print("done")
    return(apparatus)

def update_webpage(byWhom=None, apparatus=None):
    timestr = time.strftime("%a, %d %b %Y, %H:%M:%S", time.localtime())
    clock_upd.text = f'Last update at {timestr}' 
    try:
        status_line.text = apparatus.state
    except:
        pass
    jp.run_task(wp.update())

async def update_webpage_loop(update_interval=1):
    while True:
        update_webpage()
        await asyncio.sleep(update_interval)

async def getPage():
    return wp

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

if __name__ == '__main__': 
    wp = jp.WebPage(delete_flag=False)
    apparatus = BasicInstrument()

    d=jp.Div(text='Magnetometer log', a=wp, classes='text-white bg-blue-500 text-center text-xl')

    div_status = jp.Div(classes='text-xl flex m0 p-1 space-x-4 bg-gray-300 font-mono', a=wp)
    clock_upd = jp.Div(text='Clock Loading...', classes='text-xl bg-gray-400', a=div_status)
    status_line = jp.Div(text='Status Loading...',classes='text-xl', a=div_status)

    # mpl.use("TkAgg")
    # apparatus = asyncio.run(main())
    jp.justpy(getPage, startup=jp_startup)