%load_ext autoreload
%autoreload 2
import warnings
warnings.filterwarnings("ignore")

from qiskit_metal import designs, MetalGUI
from qiskit_metal.qlibrary.sample_shapes.rectangle import Rectangle
from qiskit_metal.qlibrary.qubits.transmon_pocket import TransmonPocket
from qiskit_metal.qlibrary.tlines.straight_path import RouteStraight
from qiskit_metal.qlibrary.tlines.meandered import RouteMeander
from qiskit_metal.qlibrary.terminations.open_to_ground import OpenToGround
from qiskit_metal.qlibrary.terminations.short_to_ground import ShortToGround

design = designs.MultiPlanar({}, overwrite_enabled=True, layer_stack_filename="layer_stack.txt")
design.variables.sample_holder_top = '320 um'
design.variables.sample_holder_bottom ='250 um'
gui = MetalGUI(design)

design.ls.ls_df

design.delete_all_components()
# Define a dictionary for connection pad options for the transmon
conn_pads1 = dict(connection_pads = dict(coupler = dict(loc_W=1, loc_H=1),
                                         readout = dict(loc_W=-1, loc_H=-1)))
conn_pads2 = dict(pad_width = '550 um', 
                  pad_gap = '15um', 
                  connection_pads = dict(coupler = dict(loc_W=-1, loc_H=1,pad_width='275um', pad_cpw_shift = '12.5um'),
                                         readout = dict(loc_W=1, loc_H=-1,pad_width='275um')))

# Create a TransmonPocket6 object
q1 = TransmonPocket(design, "Q1", options=dict(pos_x='-0.5mm', pos_y='+0.0mm', layer=1, **conn_pads1))
q2 = TransmonPocket(design, "Q2", options=dict(pos_x='+0.5mm', pos_y='+0.0mm', layer=1, **conn_pads2))

# Rebuild and autoscale the GUI
gui.rebuild()
gui.autoscale()

gui.screenshot()

# Connect Transmons with direct coupler
coupler_options = dict(pin_inputs = dict(start_pin=dict(component=q2.name, pin='coupler'),
                                         end_pin=dict(component=q1.name, pin='coupler')))

coupler_connector = RouteStraight(design, 'coupler_connector', options= coupler_options)

# Rebuild and autoscale the GUI
gui.rebuild()
gui.autoscale()

gui.screenshot()

# Create open-to-ground elements for readout resonators
otg11_options = dict(pos_x='-1.5mm',  pos_y='0.0mm', orientation='180', layer='3')
otg21_options = dict(pos_x='+1.5mm',  pos_y='0.0mm', orientation='0', layer='3')
otg12_options = dict(pos_x='-0.955mm',  pos_y='-0.195mm', orientation='0', layer='3')
otg22_options = dict(pos_x='+0.955mm',  pos_y='-0.1875mm', orientation='180', layer='3')
otg11 = OpenToGround(design, 'otg11', options=otg11_options)
otg21 = OpenToGround(design, 'otg21', options=otg21_options)
otg12 = OpenToGround(design, 'otg12', options=otg12_options)
otg22 = OpenToGround(design, 'otg22', options=otg22_options)


# Create readout resonators
readout1_options = dict(total_length = '5.97mm',
                       fillet = '40um',
                       pin_inputs = dict(start_pin = dict(component = otg11.name, pin = 'open'),
                                         end_pin = dict(component = otg12.name, pin = 'open')),
                       lead=dict(start_straight='150um'),
                       meander=dict(spacing = '100um', asymmetry = '0'),
                       layer = '3')

readout2_options = dict(total_length = '5.97mm',
                       fillet = '40um',
                       pin_inputs = dict(start_pin = dict(component = otg22.name, pin = 'open'),
                                         end_pin = dict(component = otg21.name, pin = 'open')),
                       lead=dict(start_straight='120um'),
                       meander=dict(spacing = '100um', asymmetry = '0.2'),
                       layer = '3')

readout1 = RouteMeander(design, 'readout1', options=readout1_options)
readout2 = RouteMeander(design, 'readout2', options=readout2_options)

# Rebuild and autoscale the GUI
gui.rebuild()
gui.autoscale()

gui.screenshot()

# Create via components
via_positions = [('-0.94mm','-0.1950mm'),('+0.94mm','-0.1875mm')]

for layer in ['1','2','3']:
    for via, positions in enumerate(via_positions):
        for subtract in zip([False,True],['30um','42um']):
            via_size = '20um' if layer == '2' else subtract[1]
            actual_layer = '5' if (layer == '2' and subtract[0]) else layer
            via_options = dict(width = via_size,
                               height = via_size,
                               pos_x = positions[0],
                               pos_y = positions[1],
                               subtract = subtract[0],
                               layer = actual_layer,
                               orientation = '0',
                               helper = 'False',
                               chip = 'main')

            name = 'via'+str(via+1)+'_layer'+actual_layer+('' if not subtract[0] else '_sub')
            Rectangle(design, name, via_options)

# Rebuild and autoscale the GUI
gui.rebuild()
gui.autoscale()

gui.screenshot()

from qiskit_metal.renderers.renderer_elmer.elmer_renderer import QElmerRenderer

elmer_renderer = QElmerRenderer(design, layer_types=dict(metal=[1,2,3], dielectric=[4,5]))
elmer_renderer.render_design(mesh_geoms=True, 
                             skip_junctions=True,
                             open_pins=[("Q1", "readout"),("Q2", "readout")],
                             omit_ground_for_layers=[2])

#elmer_renderer.launch_gmsh_gui()

elmer_renderer.export_mesh()
elmer_renderer.add_solution_setup('capacitance')
elmer_renderer.run('capacitance')

# Display capacitance matrix in [fF]
cap_matrix = elmer_renderer.capacitance_matrix
cap_matrix

elmer_renderer.display_post_processing_data()
elmer_renderer.close()

# Mesh Q1-only design
elmer_renderer = QElmerRenderer(design, layer_types=dict(metal=[1,2,3], dielectric=[4,5]))
select = ['Q1', 'via1_layer1', 'via1_layer2', 'via1_layer3',
          'via1_layer1_sub', 'via1_layer5_sub', 'via1_layer3_sub']

elmer_renderer.render_design(selection=select,
                             open_pins=[('Q1', 'coupler'), ('Q1', 'readout')],
                             skip_junctions=True,
                             omit_ground_for_layers=[2])

elmer_renderer.launch_gmsh_gui()

# Solve Q1-only mode and save Cap Matrix
elmer_renderer.export_mesh()
elmer_renderer.add_solution_setup('capacitance')
elmer_renderer.run('capacitance')
cap_matrix_q1 = elmer_renderer.capacitance_matrix
#elmer_renderer.close()
cap_matrix_q1

elmer_renderer.display_post_processing_data()
elmer_renderer.close()

# Mesh Q2-only design
elmer_renderer = QElmerRenderer(design, layer_types=dict(metal=[1,2,3], dielectric=[4,5]))
select = ['Q2', 'via2_layer1', 'via2_layer2', 'via2_layer3',
          'via2_layer1_sub', 'via2_layer5_sub', 'via2_layer3_sub']
elmer_renderer.render_design(selection=select,
                             open_pins=[('Q2', 'coupler'), ('Q2', 'readout')],
                             skip_junctions=True, omit_ground_for_layers=[2])
#elmer_renderer.launch_gmsh_gui()

# Solve Q2-only mode and save Cap Matrix
elmer_renderer.export_mesh()
elmer_renderer.add_solution_setup('capacitance')
elmer_renderer.run('capacitance')
cap_matrix_q2 = elmer_renderer.capacitance_matrix
elmer_renderer.close()
cap_matrix_q2

%matplotlib inline
import scqubits as scq
from scipy.constants import speed_of_light as c_light
import matplotlib.pyplot as plt

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

# cell 1: 1st transmon (Q1)
opt1 = dict(node_rename = {'Q1_coupler_connector_pad': 'coupling', 'via1_layer3_rectangle': 'readout_1'}, 
            cap_mat = cap_matrix_q1,
            ind_dict = {('Q1_pad_top', 'Q1_pad_bot'):10},  # junction inductance in nH
            jj_dict = {('Q1_pad_top', 'Q1_pad_bot'):'j1'},
            cj_dict = {('Q1_pad_top', 'Q1_pad_bot'):2} # junction capacitance in fF
           )
cell_1 = Cell(opt1)


# cell 2: 2nd transmon (Q2)
opt2 = dict(node_rename = {'Q2_coupler_connector_pad': 'coupling', 'via2_layer3_rectangle': 'readout_2'},
            cap_mat = cap_matrix_q2,
            ind_dict = {('Q2_pad_top', 'Q2_pad_bot'): 12}, # junction inductance in nH
            jj_dict = {('Q2_pad_top', 'Q2_pad_bot'):'j2'},
            cj_dict = {('Q2_pad_top', 'Q2_pad_bot'):2} # junction capacitance in fF
           )
cell_2 = Cell(opt2)

# subsystem 1: transmon 1
transmon_1 = Subsystem(name='transmon_1', sys_type='TRANSMON', nodes=['j1'])

# subsystem 2: transmon 2
transmon_2 = Subsystem(name='transmon_2', sys_type='TRANSMON', nodes=['j2'])

# subsystem 3: readout resonator 1
q_opts = dict(f_res = 8, # resonator dressed frequency in GHz
              Z0 = 50, # characteristic impedance in Ohm
              vp = 0.404314 * c_light # phase velocity 
              )
res_1 = Subsystem(name='readout_1', sys_type='TL_RESONATOR', nodes=['readout_1'], q_opts=q_opts)

# subsystem 4: readout resonator 2
q_opts = dict(f_res = 7.6, # resonator dressed frequency in GHz
              Z0 = 50, # characteristic impedance in Ohm
              vp = 0.404314 * c_light # phase velocity 
              )
res_2 = Subsystem(name='readout_2', sys_type='TL_RESONATOR', nodes=['readout_2'], q_opts=q_opts)

composite_sys = CompositeSystem(subsystems=[transmon_1, transmon_2, res_1, res_2], 
                                cells=[cell_1, cell_2], 
                                grd_node='ground_plane',
                                nodes_force_keep=['readout_1', 'readout_2'])

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

