%load_ext autoreload
%autoreload 2


import scqubits as scq

from qiskit_metal.analyses.quantization.lumped_capacitive import load_q3d_capacitance_matrix
from qiskit_metal.analyses.quantization.lom_core_analysis import CompositeSystem, Cell, Subsystem, QuantumSystemRegistry
                        
from scipy.constants import speed_of_light as c_light

import matplotlib.pyplot as plt
%matplotlib inline

QuantumSystemRegistry.registry()

# Q1
path1 = "./capacitance_matrices/Q1_cap_matrix.txt"
t1_mat, _, _, _ = load_q3d_capacitance_matrix(path1)

# Q2
path2 = "./capacitance_matrices/Q2_cap_matrix.txt"
t2_mat, _, _, _ = load_q3d_capacitance_matrix(path2)

# Q3
path3 = "./capacitance_matrices/Q3_cap_matrix.txt"
t3_mat, _, _, _ = load_q3d_capacitance_matrix(path3)

# Q4
path4 = "./capacitance_matrices/Q4_cap_matrix.txt"
t4_mat, _, _, _ = load_q3d_capacitance_matrix(path4)

# Cell 2: Transmon-2
opt2 = dict(
    node_rename = {'b_connector_pad_Q2': 'coupler23',
                   'c_connector_pad_Q2': 'readout2'}, 
    cap_mat = t2_mat,
    ind_dict = {('pad_top_Q2', 'pad_bot_Q2'): 12},  # junction inductance in nH
    jj_dict = {('pad_top_Q2', 'pad_bot_Q2'): 'j2'},
    cj_dict = {('pad_top_Q2', 'pad_bot_Q2'): 1}, # junction capacitance in fF

)
cell_2 = Cell(opt2)

# Cell 3: Transmon-3
opt3 = dict(
    node_rename = {'b_connector_pad_Q3': 'coupler23',
                   'c_connector_pad_Q3': 'readout3'}, 
    cap_mat = t3_mat,
    ind_dict = {('pad_top_Q3', 'pad_bot_Q3'): 12},  # junction inductance in nH
    jj_dict = {('pad_top_Q3', 'pad_bot_Q3'): 'j3'},
    cj_dict = {('pad_top_Q3', 'pad_bot_Q3'): 1}, # junction capacitance in fF

)
cell_3 = Cell(opt3)

# subsystem 1: Transmon-2
transmon2 = Subsystem(name='transmon2', sys_type='TRANSMON', nodes=['j2'])

# subsystem 2: Transmon-3
transmon3 = Subsystem(name='transmon3', sys_type='TRANSMON', nodes=['j3'])

# Resonator Subsystems
q_opts = dict(
    Z0 = 50, # characteristic impedance in Ohm
    vp = 0.404314 * c_light # phase velocity 
)

# subsystem 3: Q2 readout resonator
ro2 = Subsystem(name='readout2', sys_type='TL_RESONATOR', nodes=['readout2'], q_opts=dict(f_res = 6.96, **q_opts))

# subsystem 4: Q3 readout resonator
ro3 = Subsystem(name='readout3', sys_type='TL_RESONATOR', nodes=['readout3'], q_opts=dict(f_res = 6.98, **q_opts))

# subsystem 15: Q2-Q3 coupling resonator
coup23 = Subsystem(name='coupler23', sys_type='TL_RESONATOR', nodes=['coupler23'], q_opts=dict(f_res = 7.36, **q_opts))

composite_sys = CompositeSystem(
    subsystems=[transmon2, transmon3, coup23, ro2, ro3], 
    cells=[cell_2, cell_3],
    grd_node='ground_main_plane',
    nodes_force_keep=['readout2', 'readout3', 'coupler23']
)

cg = composite_sys.circuitGraph()
print(cg)

hilbertspace = composite_sys.create_hilbertspace()
print(hilbertspace)

# Convert the hilbert space into
# "Interaction Picture"
hilbertspace = composite_sys.add_interaction()
hilbertspace.hamiltonian()

hamiltonian_results = composite_sys.hamiltonian_results(hilbertspace, evals_count=30)

composite_sys.compute_gs()

