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 pyvisa
import numpy as np
import platform
import time
from tqdm import tqdm as pbar


from qolab.hardware import BasicInstrument
from qolab.data import TraceSetSameX, TraceXY, Trace
from qolab.hardware.rf_generator import QOL_LMX2487
from qolab.hardware.daq import LabJackUE9
from qolab.hardware.power_supply.keysight_e3612a import KeysightE3612A
import qolab.tsdb as tsdb

l = logging.getLogger('qolab.tsdb')
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

class Apparatus(BasicInstrument):
    def __init__(self, *args, **kwds):
        super().__init__(*args, **kwds)

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

# TSDB logger setting
tsdb_client = tsdb.Client('influx', 'http://lumus.physics.wm.edu:8428', database='qolab')
tsdb_ingester = tsdb.Ingester(tsdb_client, batch=11, measurement_prefix='VAMPIRE.VICSEL')

# creating Apparatus with all instruments to be logged
app_nickname = 'magnetometer'
apparatus = Apparatus(tsdb_ingester=tsdb_ingester, device_nickname=app_nickname,)
apparatus.config['Device type'] = 'QOL VAMPIRE VICSEL magnetometer'
apparatus.config['Device model'] = 'v0.1'
apparatus.config['FnamePrefix'] = 'magnetometer_eit'
# apparatus.config['SavePath'] = '/mnt/qol_grp_data/data.VAMPIRE'
apparatus.config['SavePath'] = './data'


logger.info("Accessing hardware")
rm = pyvisa.ResourceManager()
instr=rm.open_resource('USB0::10893::4354::MY61001869::0::INSTR')
ps = KeysightE3612A(instr, device_nickname='.'.join([app_nickname, 'coil_driver']), tsdb_ingester=tsdb_ingester)

if platform.system() == 'Linux':
    rfgen_port='/dev/ttyUSB0'
else:
    rfgen_port='COM4'
rfgen = QOL_LMX2487(port=rfgen_port, speed=115200, timeout=1, device_nickname='.'.join([app_nickname, 'rfgen']), tsdb_ingester=tsdb_ingester)


daq = LabJackUE9(device_nickname='.'.join([app_nickname, 'daq']), tsdb_ingester=tsdb_ingester)

logger.info("Adding instruments to apparatus")
apparatus.instruments={}
ai = apparatus.instruments
ai['rfgen'] = rfgen
ai['daq']   = daq
ai['coil_driver'] = ps


logger.info('Setting magnetic field coils currents')
"""
Rough magnetic field calibration of the 3 axes coils
- Ch1: 70mA -> 650 kHz shift for delta m = 2
- Ch2: 70mA -> 700 kHz shift for delta m = 2
- Ch2: 70mA -> 659 kHz shift for delta m = 2
"""
ps.setChanIout_mA(1, 70)
ps.setChanIout_mA(2, 0)
ps.setChanIout_mA(3, 0)
logger.info('Done setting magnetic field coils currents')

def eitSweep(central_frequency, frequency_span, Np, Nsweeps=1):
    frList = np.linspace(central_frequency-frequency_span/2, central_frequency+frequency_span/2, Np)
    trFreq=Trace('Frequency')
    trFreq.config['unit']='Hz'
    trTransmission=Trace('Transmission')
    trTransmission.config['unit']='Arb. Unit'
    trLockin=Trace('Lockin')
    trLockin.config['unit']='V'

    trEIT = TraceSetSameX('EIT')

    for sw in pbar(range(1,Nsweeps+1), desc="Sweep"):
        for fr in pbar(frList, desc="Freq Scan"):
            rfgen.setFreqFixed(float(fr))
            time.sleep(dwellTime)
            transmission = daq.getAIN(0)
            lockin = daq.getAIN(1)

            trFreq.addPoint(fr)
            trTransmission.addPoint(transmission)
            trLockin.addPoint(lockin)

    # trFreq.values = trFreq.values - central_frequency
    trEIT.addTraceX(trFreq)
    trEIT.addTrace(trTransmission)
    trEIT.addTrace(trLockin)

    trEIT.config['tags']['apparatus']=apparatus.getConfig()

    return trEIT

def eitVsCurrent(ch=1):
    curList = [20, 30, 40, 50, 60, 70, 80, 90, 100, 110]
    ps.setChanIout_mA(1, 0)
    ps.setChanIout_mA(2, 0)
    ps.setChanIout_mA(3, 0)
    for current in curList:
        logger.info(f'Preparing data for {current=} in {ch=} ready')
        ps.setChanIout_mA(ch, current)
        trEIT = eitSweep(central_frequency, frequency_span, Np, Nsweeps=Nsweeps)
        trEIT.plot()
        fn = apparatus.getNextDataFile()
        logger.info(f'Data ready for {current=} in {ch=}')
        logger.info(f'Data saved to {fn=}')
        trEIT.save(fn)

central_frequency = 6.83468e9
frequency_span =  2500e3

dwellTime=0.1
Np=1000
Nsweeps=5


eitVsCurrent(ch=1)
eitVsCurrent(ch=2)
eitVsCurrent(ch=3)

tsdb_ingester.commit()