import logging
# this should be done before justpy is called or log formatter does not work
logging.basicConfig(format='%(asctime)s %(levelname)8s %(name)s: %(message)s', datefmt='%m/%d/%Y %H:%M:%S')

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.hardware.lockin import SRS_SR865A
from qolab.feedback import PID
from qolab.data import TraceSetSameX, TraceXY, Trace

l = logging.getLogger('qolab.gui.web')
l.setLevel(logging.INFO)
logger = logging.getLogger('Magnetometer')
logger.setLevel(logging.INFO)

def getConfig(apparatus):
    config = apparatus.config.copy()
    ai = apparatus.instruments
    for n, i in ai.items():
        config[n]=i.getConfig()
    return config

def initLog():
    errorTrace = Trace("error")
    errorTrace.config['unit']='V'
    timeTrace = Trace("time")
    timeTrace.config['unit']='S'
    timeTrace.config['type'] = 'timestamp'
    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

    eitTrace = Trace("eit")
    eitTrace.config['unit']='V'
    eitLog = TraceXY("eit")
    eitLog.x = timeTrace
    eitLog.y = eitTrace

    rfPoutTrace = Trace("rfPout")
    rfPoutTrace.config['unit']='dBm'
    rfPoutLog = TraceXY("rfPout")
    rfPoutLog.x = timeTrace
    rfPoutLog.y = rfPoutTrace


    log = TraceSetSameX("timelog")
    log.addTrace(errorLog)
    log.addTrace(freqLog)
    log.addTrace(feedbackLog)
    log.addTrace(eitLog)
    log.addTrace(rfPoutLog)
    log.config['tags']['apparatus']=getConfig(apparatus)
    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()

    ai = apparatus.instruments
    scope = ai['scope']
    rfgen = ai['rfgen']
    pid   = ai['pid']
    log   = apparatus.gui_log.traces

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

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

    rfPout = ai['rfgen'].getRFAmplitude()

    log.addPointToTrace(timenow)
    log.addPointToTrace(err, "error")
    log.addPointToTrace(freq0, "frequency")
    log.addPointToTrace(fdbck, "feedback")
    log.addPointToTrace(rfPout, "rfPout")
    log.addPointToTrace(eit, "eit")


freqZero = 6834686400
freqDeltaMp2 = 6835387100
dfB = freqDeltaMp2 - freqZero
freqDeltaMm2 = freqZero - dfB;

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

    logger.info("Accessing hardware")
    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)
    instr_lockin=rm.open_resource('TCPIP::192.168.0.51::INSTR')
    lockin = SRS_SR865A(instr_lockin)
    # 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); # goom for dm=-2 resonance
    pid = PID(50,200,0, sign=-1)

    apparatus.instruments={}
    ai = apparatus.instruments
    ai['rfgen'] = rfgen
    ai['lockin']   = lockin
    ai['pid']   = pid
    ai['scope'] = scope
    apparatus.state = None
    apparatus.runStatus = False


    SweepSpan = 10000
    ai['rfgen'].stopFrequencySweep()
    # ai['rfgen'].setFreqFixed(freqDeltaMm2)
    # ai['rfgen'].setFreqFixed(freqZero)
    ai['rfgen'].setFreqFixed(freqDeltaMp2)

    rfPstart = -10 
    rfPstop = 10
    rfPowerList = np.linspace(rfPstart, rfPstop, 11)
    rfPower0 = rfPstart
    rfPower0 = 3.4 
    ai['rfgen'].setRFAmplitude(rfPower0)

    d=getConfig(apparatus)
    instruments_config=gui.QOLDictionary(a=wp, name='Instruments configs', container=d)

    logger.info('========== Initial lock ===========')
    rfPout = ai['rfgen'].getRFAmplitude()
    apparatus.state = f'Initial lock RF power {rfPout} dBm'
    logger.info(apparatus.state)
    update_webpage(apparatus=apparatus)
    apparatus.gui_log.setTraces( initLog() )
    ai['pid'].reset()
    ai['pid'].setEnable(True)
    apparatus.runStatus = True
    res = await asyncio.gather(
            feedbackLoop(apparatus, nsteps=200)
            )

    async def calibratingLockin():
        apparatus.state = f'Calibrating lockin response at RF power {rfPout} dBm'
        logger.info(apparatus.state)
        update_webpage(apparatus=apparatus)
        ai['pid'].setEnable(False)
        ai['pid'].reset()
        apparatus.gui_log.setTraces( initLog() )
        res = await asyncio.gather(
                feedbackLoop(apparatus, nsteps=20)
                )

        fr0 = ai['rfgen'].getFreqFixed()
        df = 10
        ai['pid'].setEnable(False)
        ai['pid'].reset()
        ai['rfgen'].setFreqFixed(fr0+df)
        res = await asyncio.gather(
                feedbackLoop(apparatus, nsteps=20)
                )
        ai['rfgen'].setFreqFixed(fr0)

        log   = apparatus.gui_log.traces
        trE=log.getTrace('error')

        e1 = np.mean(trE.y.values[0:20])
        e2 = np.mean(trE.y.values[20:])
        dE = e2-e1
        slope = dE/df

        logger.info(f'dE = {dE}')
        logger.info(f'slope = {slope}')

        return slope

    lockin_slope= await calibratingLockin()

    logger.info('========= Long term lock ===========')
    apparatus.gui_log.setTraces( initLog() )
    apparatus.gui_log.traces.config['tags']['lockin_slope']=float(lockin_slope)
    apparatus.runStatus = True

    async def longTermLock():
        apparatus.state = f'Long term lock RF power {rfPout} dBm'
        ai['pid'].setEnable(True)
        logger.info(apparatus.state)
        update_webpage(apparatus=apparatus)
        res = await asyncio.gather(
                feedbackLoop(apparatus, nsteps=100000)
                )
        return res

    async def sweepRFPower():
        # for p in np.linspace(-10,10, 110):
        ai['pid'].setEnable(True)
        for p in rfPowerList:
            ai['rfgen'].setRFAmplitude(p)
            rfPout = ai['rfgen'].getRFAmplitude()
            apparatus.state = f'Long term lock RF power {rfPout} dBm'
            logger.info(apparatus.state)
            update_webpage(apparatus=apparatus)
            ai['pid'].reset()
            apparatus.runStatus = True
            res = await asyncio.gather(
                    feedbackLoop(apparatus, nsteps=50)
                    )
        return res

    # await sweepRFPower()
    await longTermLock()
    apparatus.state = 'Done working with hardware'
    logger.info(apparatus.state)

    apparatus.gui_log.stop_tasks()
    task_wp_update_loop.cancel()
    logger.info("exiting main loop")
    update_webpage(apparatus=apparatus)
    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__': 
    logger.info('========== Start up ===========')
    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)