import pyvisa
import numpy as np
import platform
import time
from tqdm import tqdm as pbar


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

rm = pyvisa.ResourceManager()
instr=rm.open_resource('USB0::10893::4354::MY61001869::0::INSTR')
ps = KeysightE3612A(instr)

if platform.system() == 'Linux':
    rfgen=QOL_LMX2487(port='/dev/ttyUSB0', speed=115200, timeout=1)
else:
    rfgen=QOL_LMX2487(port='COM4', speed=115200, timeout=1)

daq = LabJackUE9()

central_frequency = 6.83468e9
frequency_span =  2500e3

dwellTime=0.1
Np=100

print('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)
print('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(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)

    return trEIT

trEIT = eitSweep(central_frequency, frequency_span, Np)
trEIT.plot()