%load_ext autoreload
%autoreload 2

import os
import warnings
import numpy as np
import pandas as pd
import qiskit_metal as metal
from numpy import pi
from scipy.constants import c, h, pi, hbar, e
from qiskit_metal import designs, draw, MetalGUI, Dict, Headings
from qiskit_metal.qlibrary.qubits.transmon_pocket_6 import TransmonPocket6
from qiskit_metal.qlibrary.qubits.transmon_pocket_cl import TransmonPocketCL
from qiskit_metal.qlibrary.tlines.meandered import RouteMeander
from qiskit_metal.qlibrary.terminations.open_to_ground import OpenToGround
from qiskit_metal.qlibrary.lumped.cap_3_interdigital import Cap3Interdigital
from qiskit_metal.qlibrary.terminations.launchpad_wb_coupled import LaunchpadWirebondCoupled
from qiskit_metal.analyses.em.cpw_calculations import guided_wavelength
from qiskit_metal.analyses.quantization import LOManalysis, EPRanalysis
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")

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

#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, 9, 4),2), 10, 6, 2)

chipX = 9
chipY = 9

design.chips.main.size_x = str(chipX)+"mm"
design.chips.main.size_y = str(chipY)+"mm"

gui = MetalGUI(design)

design.delete_all_components()

# basic parameters for inital cut
C_JJ = 2

pad_g = 30
pad_w = 450
pad_h = 150

readout_g  = 30
buslineL_g = 30
buslineH_g = 30

C_JJU = str(C_JJ)+"fF"
pad_g = str(pad_g)+"um"
pad_w = str(pad_w)+"um"
pad_h = str(pad_h)+"um"

readout_g  = str(readout_g)+"um"
buslineL_g = str(buslineL_g)+"um"
buslineH_g = str(buslineH_g)+"um"

connector1_pads_options = dict(
    cR = dict(loc_W=0, loc_H=1, pad_width='80 um', pad_gap=readout_g, pad_height='30 um'),  # readout resonator
    cL = dict(loc_W=1, loc_H=-1, pad_width='80 um', pad_gap=buslineL_g, pad_height='30 um'), # bus line for lower frequency
    cH = dict(loc_W=-1, loc_H=-1, pad_width='80 um', pad_gap=buslineH_g, pad_height='30 um') # bus line for higher frequency
)
connector2_pads_options = dict(
    cR = dict(loc_W=0, loc_H=1, pad_width='80 um', pad_gap=readout_g, pad_height='30 um'),  # readout resonator
    cL = dict(loc_W=1, loc_H=-1, pad_width='80 um', pad_gap=buslineL_g, pad_height='30 um'), # bus line for lower frequency
    cH = dict(loc_W=-1, loc_H=-1, pad_width='80 um', pad_gap=buslineH_g, pad_height='30 um') # bus line for higher frequency
)
connector3_pads_options = dict(
    cR = dict(loc_W=0, loc_H=1, pad_width='80 um', pad_gap=readout_g, pad_height='30 um'),  # readout resonator
    cL = dict(loc_W=1, loc_H=-1, pad_width='80 um', pad_gap=buslineL_g, pad_height='30 um'), # bus line for lower frequency
    cH = dict(loc_W=-1, loc_H=-1, pad_width='80 um', pad_gap=buslineH_g, pad_height='30 um') # bus line for higher frequency
)
connector4_pads_options = dict(
    cR = dict(loc_W=0, loc_H=1, pad_width='80 um', pad_gap=readout_g, pad_height='30 um'),  # readout resonator
    cL = dict(loc_W=1, loc_H=-1, pad_width='80 um', pad_gap=buslineL_g, pad_height='30 um'), # bus line for lower frequency
    cH = dict(loc_W=-1, loc_H=-1, pad_width='80 um', pad_gap=buslineH_g, pad_height='30 um') # bus line for higher frequency
)

# target f: 5.25 
connection1_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)
)

Q1 = TransmonPocketCL(design, 'Q1',
                      options=dict(pos_x=f'0mm',
                                   pos_y=f'0mm',
                                   pad_gap = pad_g,
                                   pad_width = pad_w,
                                   pad_height = pad_h,
                                   orientation=0,
                                   connection_pads=dict(**connection1_pads_options)))

gui.rebuild()
gui.autoscale()

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

hfss.activate_ansys_design("Q1Readout", 'eigenmode')

hfss.render_design(['Q1'], [])

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

pinfo = hfss.pinfo
pinfo.design.set_variable('Lj', str(L_JJList[0])+'nH')
pinfo.design.set_variable('Cj', C_JJU)

setup.analyze()

eig_qres.sim.convergence_t, eig_qres.sim.convergence_f, _ = hfss.get_convergences()
eig_qres.sim.plot_convergences()

import pyEPR as epr

pinfo = hfss.pinfo
pinfo.junctions['jj'] = {'Lj_variable': 'Lj', 'rect': 'JJ_rect_Lj_Q1_rect_jj',
                             'line': 'JJ_Lj_Q1_rect_jj_',  'Cj_variable': 'Cj'}
pinfo.validate_junction_info() # Check that valid names of variables and objects have been supplied
pinfo.dissipative['dielectrics_bulk'] = ['main'] # Dissipative elements: specify

eprd = epr.DistributedAnalysis(pinfo)

print(eprd)

ℰ_elec = eprd.calc_energy_electric()
ℰ_elec_substrate = eprd.calc_energy_electric(None, 'main')
ℰ_mag = eprd.calc_energy_magnetic()

print(f"""
ℰ_elec_all       = {ℰ_elec}
ℰ_elec_substrate = {ℰ_elec_substrate}
EPR of substrate = {ℰ_elec_substrate / ℰ_elec * 100 :.1f}%

ℰ_mag_all       = {ℰ_mag}
ℰ_mag % of ℰ_elec_all  = {ℰ_mag / ℰ_elec * 100 :.1f}%
""")

eprd.do_EPR_analysis()

# 4a. Perform Hamiltonian spectrum post-analysis, building on mw solutions using EPR
epra = epr.QuantumAnalysis(eprd.data_filename)
epra.analyze_all_variations(cos_trunc = 8, fock_trunc = 7)

hfss.modeler._modeler.ShowWindow()
hfss.plot_fields('main')