from entropylab.graph_experiment import Graph
from entropylab.instruments.lab_topology import ExperimentResources
from entropylab.results_backend.sqlalchemy.db import SqlAlchemyDB
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from qpu_resolver import resolve
# from nodes import *
from test_root import root
from test_q_spec import QubitSpecNode
import os, qm, sys

sys.path.append(r"C:\Program Files\Labber\Script")
import Labber
from labber_helper_functions import *


# connect to entropy db
working_dir = os.getcwd()
db_path = os.path.join(os.path.dirname(os.getcwd()), 'database')
entropy_db_file = os.path.join(db_path, "entropy_db_20220601_152358.db")
entropy_db = SqlAlchemyDB(entropy_db_file)

# define resources needed for the experiment
experiment_resources = ExperimentResources(entropy_db)
experiment_resources.import_lab_resource('qmm')
experiment_resources.import_lab_resource('qpu_db')
experiment_resources.pause_save_to_results_db()
## register Labber instruments
Labber_Instruments_list = ['XY_LO', 'RO_LO', 'TWPA1_LO', 'TWPA2_LO', 'QDAC', 'LabBrick_RO', 'RSA', 
                           'LabBrick_DC1', 'LabBrick_DC2', 'BRF1', 'BRF2', 'MixerCalSwitch']
for j, ins in enumerate(Labber_Instruments_list):
    experiment_resources.import_lab_resource(ins)

mixer_cal_switch = experiment_resources.get_resource('MixerCalSwitch')
# mixer_cal_switch.switch_all_to_experiment_mode()

qpu_db = experiment_resources.get_resource('qpu_db')


# %% import crosstalk matrix from qpu_db

num_qubits = qpu_db.num_qubits

ct_mat = np.array([
    [qpu_db.get('crosstalk_mat_dc', f'C_{k1 + 1},{k2 + 1}').value
     for k2 in range(num_qubits)]
    for k1 in range(num_qubits)])
inv_ct_mat = np.linalg.inv(ct_mat)

# % initial conditions from previous characterizations
# bias voltages of qubits (with other qubits at 0V) used for flux crosstalk calibration
ctMeas_vol_list = [-2.95, -3.792, -3.226,
                   -4.123, -3.53, -3.85,
                   -3.35, -4.05, -3.1,
                   -4.25]

V_ideal_ctMeas = np.zeros(num_qubits)
for i in range(num_qubits):
    V_appl = np.zeros(num_qubits)
    V_appl[i] = ctMeas_vol_list[i]

    V_ideal_ctMeas[i] = np.matmul(ct_mat, V_appl)[i]

print("Ideal voltages:", V_ideal_ctMeas)

readout_freqs = [5.846e9, 6.112e9, 6.363e9, 5.5930e9, 6.0375e9,
                 5.923e9, 5.683e9, 6.244e9, 6.187e9, 5.742e9]  ## lower bandgap

# readout_freqs = [5.79e9, 6.03e9, 6.27e9, 5.54e9, 5.953e9,
#                  5.851e9, 5.613e9, 6.18e9, 6.187e9, 5.742e9]
print("Readout frequencies", readout_freqs)

qubit_freqs = [4.501e9, 4.5350e9, 4.5070e9, 4.1e9, 4.488e9,
               4.489e9, 4.521e9, 4.442e9, 4.413e9, 4.375e9]

print("Qubit frequencies", qubit_freqs)
# %%

qpu_db = experiment_resources.get_resource('qpu_db')

# % import crosstalk matrix from qpu_db

num_qubits = qpu_db.num_qubits

ct_mat = np.array([
    [qpu_db.get('crosstalk_mat_dc', f'C_{k1 + 1},{k2 + 1}').value
     for k2 in range(num_qubits)]
    for k1 in range(num_qubits)])
inv_ct_mat = np.linalg.inv(ct_mat)


f_sweet_list = [3737417242.0, 3520000000.0, 3599473214.285714,
                3512000000.0, 3724000000.0, 3450900000.0, 
                3672200000.0, 3529200963.71885, 3555812563.2484784, 
                3394000000.0]

f_sweet_list_measure = [3737417242.0, 3520000000.0, 3604473214.285714,
                3512000000.0, 3738500000.0, 3469900000.0, 
                3672200000.0, 3529200963.71885, 3588812563.2484784, 
                3394000000.0]

ro_f_lo = 6e9
xy_f_lo = 4.5e9

# fq_target_dict = {1: 4.35e09, 
#                   2: 4.03e09, 
#                   3: 4.460e09, 
#                   4: 4.096e09,
#                   5: 4.521e09, 
#                   6: 4.211e09, 
#                   7: 3.920e09, 
#                   8: 4.145e9,
#                   9: 4.434e9, 
#                   10: 4.004e+09}

# 4.51811704e+09, 4.14120262e+09, 4.58271412e+09, 4.10000000e+09,
#        4.17556340e+09, 4.54470440e+09, 4.21124474e+09, 4.60000000e+09,
#        4.12000000e+09, 4.51023023e+09

# target_fq_list = [4.23136181e+09, 4.05000000e+09, 4.55104606e+09, 4.60000000e+09,
#         4.51023023e+09, 4.18989500e+09, 4.14000000e+09, 4.57716178e+09,
#         4.22000000e+09, 4.10891118e+09]

target_fq_list = [4371160974.0, 4230244559.0, 4704386312.0, 
                  4282463892.0, 4777188279.0, 4648825617.0, 
                  4197488716.0, 4305972301.0, 4727112340.0, 4226166525.0]
# [4204002425.487339,
 # 4123870341.890262,
 # 4472376286.087953,
 # 3981410764.5743675,
 # 4398880542.255875,
 # 3857881220.3945303,
 # 3915476845.408693,
 # 3970573234.706229,
 # 4352611777.213919,
 # 4248389731.4683714]
# [4.38374212e+09, 3.97755406e+09, 4.50000000e+09, 4.06586943e+09,
#        3.93742562e+09, 4.34395152e+09, 4.02165610e+09, 4.46920101e+09,
#        4.42000000e+09, 3.90000000e+09]
# target_fq_list = [7.61e9, 7.506e9, 7.651e9, 7.526e9,
#        7.4e9, 7.475e9, 7.587e9, 7.426e9,
#        7.54e9, 7.46e9]

fq_target_dict = {**{(i + 1): target_fq_list[i] for i in range(num_qubits)}}
elements = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

ro_channels = ['1', '2']

elements = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]


# %% Qubit Tuning


# # Qubit Spectroscopy
qubit_spec_nodes = [
    QubitSpecNode(elements[i], span=800e6, df=2e6, plot=True, #wait_time=10000, 
                  n_avg=100,
                  dependency=root, #callback=qubit_spec_save_data,
                  name=f'q{i + 1}_spec', save_results=False)
    for i in range(num_qubits)]

qubit_tuning_experiment = Graph(
    experiment_resources,
    set().union([root], qubit_spec_nodes),
    'multi_qubit_tuning')


qubit_tuning_experiment.run()
