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 universal_tsdb as utsdb
import qolab.tsdb as tsdb
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():
    tsdb_client = tsdb.Client('influx', 'http://lumus.phusics.wm.edu:8428', database='qolab')
    tsdb_ingester = tsdb.Ingester(tsdb_client, batch=11, measurement_prefix='VAMPIRE.HighPower')

    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); # good 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=10000000)
                )
        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)