%load_ext autoreload
%autoreload 2

import os
import warnings
import qiskit_metal as metal
from numpy import pi
from scipy.constants import c, h, pi, hbar, e, mu_0, epsilon_0
from qiskit_metal import Dict
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

ws_path = os.getcwd()
warnings.filterwarnings("ignore")

#constants:
phi0    = h/(2*e)
varphi0 = phi0/(2*pi)

#project parameters:
L_JJ = 10 # nH
C_JJ = 2  # fF
f_r  = 8  # GHz
Z_r  = 50 # Ohm
v_r  = 0.404314 # relative velocity to c

#parameter extraction:
eps_eff = (1/v_r)**2  # effective eps_r
eps_r   = 2*eps_eff-1 # approximation

I_JJ = varphi0/(L_JJ*1e-9)/1e-9        # nA         
E_JJ = (varphi0)**2/(L_JJ*1e-9)/h/1e9 # GHz

#print
print(eps_r, I_JJ, E_JJ)

QuantumSystemRegistry.registry()

capfile = ws_path+"/Transmon_Readout_CapMatrix.txt"
t1_mat, _, _, _ = load_q3d_capacitance_matrix(capfile)

t1_mat

opt1 = dict(
    cap_mat = t1_mat,
    ind_dict = {('pad_top_Q1', 'pad_bot_Q1'): L_JJ},  # junction inductance in nH
    jj_dict = {('pad_top_Q1', 'pad_bot_Q1'): 'j1'},
    cj_dict = {('pad_top_Q1', 'pad_bot_Q1'): C_JJ}, # junction capacitance in fF

)
cell_1 = Cell(opt1)

transmon1 = Subsystem(name='transmon1', sys_type='TRANSMON', nodes=['j1'])
trans_read = Subsystem(name='readout1', sys_type='TL_RESONATOR', nodes=['readout_connector_pad_Q1'], q_opts=dict(f_res = f_r, Z0=Z_r, vp=v_r*c))

composite_sys = CompositeSystem(
    subsystems=[transmon1, trans_read], 
    cells=[cell_1],
    grd_node='ground_main_plane',
    nodes_force_keep=['readout_connector_pad_Q1']
)
cg = composite_sys.circuitGraph()
print(cg)

hilbertspace = composite_sys.create_hilbertspace()
print(hilbertspace)

hilbertspace = composite_sys.add_interaction()
hilbertspace.hamiltonian()

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

composite_sys.compute_gs()

