import os
import warnings
warnings.filterwarnings('ignore')

ws_path = os.getcwd()

import numpy as np
from scipy.constants import c, h, pi, hbar, e
from qiskit_metal.analyses.em.cpw_calculations import guided_wavelength

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

# project target parameters
f_qList = np.around(np.linspace(5.25, 5.75, 4),2) # GHz
f_rList = f_qList + 1.8 # GHz
L_JJList = np.around(varphi0**2/((f_qList*1e9+300e6)**2/(8*300e6))/h*1e9, 2) # nH

# initial CPW readout lengths
def find_resonator_length(frequency, line_width, line_gap, N): 
    #frequency in GHz
    #line_width/line_gap in um
    #N -> 2 for lambda/2, 4 for lambda/4
    
    [lambdaG, etfSqrt, q] = guided_wavelength(frequency*10**9, line_width*10**-6,
                                              line_gap*10**-6, 750*10**-6, 200*10**-9)
    return str(lambdaG/N*10**3)+" mm"

find_resonator_length(f_rList, 10, 6, 2)

find_resonator_length(np.around(np.linspace(8, 8.5, 4),2), 10, 6, 2)

from qiskit_metal import designs, MetalGUI

design = designs.DesignPlanar()
design.overwrite_enabled = True

design.chips.main.size_x = '12mm'
design.chips.main.size_y = '10mm'

gui = MetalGUI(design)

from qiskit_metal.qlibrary.qubits.transmon_pocket_cl import TransmonPocketCL

design.delete_all_components()

design_span_x = 5
design_span_y = 3

half_chip_width = design_span_x / 2
half_chip_height = design_span_y / 2

connection_pads_options = dict(
    B1 = dict(loc_W=1, loc_H=-1),
    R= dict(loc_W=1, loc_H=1),
    B2 = dict(loc_W=-1, loc_H=-1)
)

transmons = []

transmons.append(TransmonPocketCL(design, 'Q1',
                      options=dict(pos_x=f'0mm',
                                   pos_y=f'{half_chip_height}mm',
                                   orientation=0,
                                   connection_pads=dict(**connection_pads_options))))

gui.rebuild()
gui.autoscale()

transmons.append(TransmonPocketCL(design, 'Q2',
                      options=dict(pos_x=f'{half_chip_width}mm',
                                   pos_y=f'0mm',
                                   orientation=-90,
                                   connection_pads=dict(**connection_pads_options))))
gui.rebuild()
gui.autoscale()

transmons.append(TransmonPocketCL(design, 'Q3',
                      options=dict(pos_x=f'{-half_chip_width}mm',
                                   pos_y=f'0mm',
                                   orientation=90,
                                   connection_pads=dict(**connection_pads_options))))
gui.rebuild()
gui.autoscale()

transmons.append(TransmonPocketCL(design, 'Q4',
                      options=dict(pos_x=f'0mm',
                                   pos_y=f'{-half_chip_height}mm',
                                   orientation=180,
                                   connection_pads=dict(**connection_pads_options))))
gui.rebuild()
gui.autoscale()

from qiskit_metal.qlibrary.tlines.meandered import RouteMeander
from qiskit_metal import Dict

fillet='99.99um'
options = Dict(
    meander=Dict(
        lead_start='0.1mm',
        lead_end='0.1mm',
        asymmetry='0 um')
)

def connect(component_name: str, component1: str, pin1: str, component2: str, pin2: str,
            length: str,
            asymmetry='0 um', start_strght='0 um', end_strght='0 um', flip=False):
    """Connect two pins with a CPW."""
    myoptions = Dict(
        pin_inputs=Dict(
            start_pin=Dict(
                component=component1,
                pin=pin1),
            end_pin=Dict(
                component=component2,
                pin=pin2)),
        lead=Dict(
            start_straight=start_strght,
            end_straight=end_strght
        ),
        total_length=length,
    fillet = '99.9um')
    myoptions.update(options)
    myoptions.meander.asymmetry = asymmetry
    myoptions.meander.lead_direction_inverted = 'true' if flip else 'false'
    return RouteMeander(design, component_name, myoptions)

cpw = []
asym_h = 200
cpw.append(connect('cpw1', 'Q1', 'B1', 'Q2', 'B2', '7.60 mm', f'-{asym_h}um', '0.1mm', '0.1mm'))
asym_h = 200
cpw.append(connect('cpw2', 'Q4', 'B2', 'Q2', 'B1', '7.45 mm', f'{asym_h}um', '0.1mm', '0.1mm'))
asym_h = 200
cpw.append(connect('cpw3', 'Q1', 'B2', 'Q3', 'B1', '7.30 mm', f'{asym_h}um', '0.1mm', '0.1mm'))
asym_h = 200
cpw.append(connect('cpw4', 'Q4', 'B1', 'Q3', 'B2', '7.15 mm', f'-{asym_h}um', '0.1mm', '0.1mm'))

gui.rebuild()
gui.autoscale()

from qiskit_metal.qlibrary.terminations.launchpad_wb_coupled import LaunchpadWirebondCoupled

readouts_lwc = []
control_lwc = []

#Readouts
readouts_lwc.append(LaunchpadWirebondCoupled(design, 'R1',
                                         options = dict(
                                         pos_x=f'3mm',
                                         pos_y=f'{half_chip_height+3}mm',
                                         orientation = 180,
                                         lead_length = '30um')))

readouts_lwc.append(LaunchpadWirebondCoupled(design, 'R2',
                                         options = dict(
                                         pos_x=f'{half_chip_width+3}mm',
                                         pos_y=f'-3mm',
                                         orientation = 90,
                                         lead_length = '30um')))

readouts_lwc.append(LaunchpadWirebondCoupled(design, 'R3',
                                         options = dict(
                                         pos_x=f'{-half_chip_width-3}mm',
                                         pos_y=f'3mm',
                                         orientation = 0,
                                         lead_length = '30um')))

readouts_lwc.append(LaunchpadWirebondCoupled(design, 'R4',
                                         options = dict(
                                         pos_x = f'-3mm',
                                         pos_y = f'{-half_chip_height-3}mm',
                                         orientation = 0,
                                         lead_length = '30um')))

gui.rebuild()
gui.autoscale()

from qiskit_metal.qlibrary.tlines.pathfinder import RoutePathfinder
from qiskit_metal.qlibrary.lumped.cap_3_interdigital import Cap3Interdigital

cap_Q1 = Cap3Interdigital(design, 'cap_Q1', options= dict(
    pos_x='2mm', pos_y=f'{half_chip_height+2}mm', orientation='180', finger_length = '40um'))

readout_lines = []

options = Dict(
    lead=Dict(
        start_straight='400um',
        end_straight='400um'),
    fillet='99.99um')

asym = 800
readout_lines.append(connect('ol1_c', 'Q1', 'R', 'cap_Q1', 'a', '8.63 mm', f'-{asym}um'))
TL_Q1 = RoutePathfinder(design, 'TL_Q1', options = dict(
                                            fillet='99um',
                                            lead=dict(end_straight='150um'),
                                            pin_inputs=Dict(
                                                start_pin=Dict(
                                                    component='R1',
                                                    pin='tie'),
                                                end_pin=Dict(
                                                    component='cap_Q1',
                                                    pin='b')
                                            )))

cap_Q2 = Cap3Interdigital(design, 'cap_Q2', options= dict(
                                         pos_x=f'{half_chip_width+2}mm', pos_y=f'-2mm', orientation='90', finger_length = '40um'))

readout_lines.append(connect('ol2', 'Q2', 'R', 'cap_Q2', 'a', '8.42 mm', f'-{asym}um'))
TL_Q2 = RoutePathfinder(design, 'TL_Q2', options = dict(
                                            fillet='99um',
                                            lead=dict(end_straight='150um'),
                                            pin_inputs=Dict(
                                                start_pin=Dict(
                                                    component='R2',
                                                    pin='tie'),
                                                end_pin=Dict(
                                                    component='cap_Q2',
                                                    pin='b')
                                            )))

cap_Q3 = Cap3Interdigital(design, 'cap_Q3', options= dict(
                                         pos_x=f'{-half_chip_width-2}mm', pos_y=f'2mm', orientation='270', finger_length = '40um'))

readout_lines.append(connect('ol3', 'Q3', 'R', 'cap_Q3', 'a', '8.25 mm', f'-{asym}um'))
TL_Q3 = RoutePathfinder(design, 'TL_Q3', options = dict(
                                            fillet='99um',
                                            lead=dict(end_straight='150um'),
                                            pin_inputs=Dict(
                                                start_pin=Dict(
                                                    component='R3',
                                                    pin='tie'),
                                                end_pin=Dict(
                                                    component='cap_Q3',
                                                    pin='b')
                                            )))

options = Dict(
    lead=Dict(
        start_straight='400um',
        end_straight='400um'),
    fillet='99.99um')

cap_Q4 = Cap3Interdigital(design, 'cap_Q4', options= dict(
                                         pos_x=f'-2mm', pos_y=f'{-half_chip_height-2}mm', orientation='0', finger_length = '40um'))

readout_lines.append(connect('ol4', 'Q4', 'R', 'cap_Q4', 'a', '8.05 mm', f'-{asym}um'))
TL_Q4 = RoutePathfinder(design, 'TL_Q4', options = dict(
                                            fillet='99um',
                                            lead=dict(end_straight='150um'),
                                            pin_inputs=Dict(
                                                start_pin=Dict(
                                                    component='R4',
                                                    pin='tie'),
                                                end_pin=Dict(
                                                    component='cap_Q4',
                                                    pin='b')
                                            )))

gui.rebuild()
gui.autoscale()

#Controls
control_lwc.append(LaunchpadWirebondCoupled(design, 'CL1',
                                         options = dict(
                                         pos_x = f'-3mm',
                                         pos_y = f'3mm',
                                         orientation=0,
                                         lead_length = '30um')))

control_lwc.append(LaunchpadWirebondCoupled(design, 'CL2',
                                         options = dict(
                                         pos_x = f'3mm',
                                         pos_y = f'3mm',
                                         orientation = -90,
                                         lead_length = '30um')))

control_lwc.append(LaunchpadWirebondCoupled(design, 'CL3',
                                         options = dict(
                                         pos_x = f'-3mm',
                                         pos_y = f'-3mm',
                                         orientation = 90,
                                         lead_length = '30um')))

control_lwc.append(LaunchpadWirebondCoupled(design, 'CL4',
                                         options = dict(
                                         pos_x = f'2mm',
                                         pos_y = f'-3mm',
                                         orientation = 90,
                                         lead_length = '30um')))

gui.rebuild()
gui.autoscale()

from qiskit_metal.qlibrary.tlines.anchored_path import RouteAnchors
from collections import OrderedDict
import numpy as np

control_lines = []

def connectRouteAnchor(name: str,
                       component1: str, pin1: str, component2: str, pin2: str,
                       anchor_points: OrderedDict) -> RouteAnchors:

    options_line_cl = dict(
        pin_inputs = dict(start_pin = dict(component = component1, pin = pin1),
                          end_pin = dict(component = component2, pin = pin2)),
        anchors = anchor_points,
        lead = dict(start_straight = '200um',
                    end_straight = '225um'),
        fillet = fillet
    )
    
    return RouteAnchors(design, name, options_line_cl)



anchors1c = OrderedDict()
anchors1c[0] = np.array([-4, -1.42])
anchors1c[1] = np.array([-4, 2])

control_lines.append(connectRouteAnchor('line_cl1', 'Q1', 'Charge_Line', 'CL1', 'tie', anchors1c))


anchors2c = OrderedDict()
anchors2c[0] = np.array([0.08, 3.25])
anchors2c[1] = np.array([4, 3.25])

control_lines.append(connectRouteAnchor('line_cl2', 'Q2', 'Charge_Line', 'CL2', 'tie', anchors2c))

anchors3c = OrderedDict()
anchors3c[0] = np.array([-0.08, -3.25])
anchors3c[1] = np.array([-4, -3.25])

control_lines.append(connectRouteAnchor('line_cl3', 'Q3', 'Charge_Line', 'CL3', 'tie', anchors3c))

anchors4c = OrderedDict()
anchors4c[0] = np.array([4, 1.42])
anchors4c[1] = np.array([4, -2])

control_lines.append(connectRouteAnchor('line_cl4', 'Q4', 'Charge_Line', 'CL4', 'tie', anchors4c))

gui.rebuild()
gui.autoscale()

from qiskit_metal.analyses.quantization import EPRanalysis
eig_res1 = EPRanalysis(design, "hfss")
hfss1 = eig_res1.sim.renderer
hfss1.start()

transmons[0].options.pad_gap    = '40um'
transmons[0].options.pad_width  = '405um'
transmons[0].options.pad_height = '90um'
gui.rebuild()

hfss1.activate_ansys_design("Tune_Q1", 'eigenmode')
hfss1.render_design(['Q1'], [('Q1', 'c'), ('Q1', 'a'),('Q1', 'b'),('Q1', 'Charge_Line')])

# Analysis properties
setup1 = hfss1.pinfo.setup
setup1.passes = 10
setup1.n_modes = 3
print(f"""
Number of eigenmodes to find             = {setup1.n_modes}
Number of simulation passes              = {setup1.passes}
Convergence freq max delta percent diff  = {setup1.delta_f}
""")

# Next 2 lines are counterinuitive, since there is no junction in this resonator.
# However, these are necessary to make pyEPR work correctly. Please do note delete
hfss1.pinfo.design.set_variable('Lj', '12 nH')
hfss1.pinfo.design.set_variable('Cj', '1 fF')

setup1.analyze()

eig_res1.sim.convergence_t, eig_res1.sim.convergence_f, _ = hfss1.get_convergences()
eig_res1.sim.plot_convergences()