%load_ext autoreload
%autoreload 2

%config IPython.sys.argv=['-Xfrozen_modules=off']

from qiskit_metal.designs.design_multiplanar import MultiPlanar

import numpy as np
import qiskit_metal as metal
from qiskit_metal import designs, draw
from qiskit_metal import MetalGUI, Dict, open_docs

design = MultiPlanar(metadata={},
                     overwrite_enabled=True)

# Design variables
design._chips.main.size.size_x = '5mm'
design._chips.main.size.size_y = '5mm'


# Check if layer data is useable, if so launch GUI
ls_unique = design.ls.is_layer_data_unique()
if (ls_unique != True):
    raise ValueError('Layer data in `MultiPlanar` design is not unique')

# gui = MetalGUI(design)

metal.__version__

# Transmon
from qiskit_metal.qlibrary.qubits.transmon_cross import TransmonCross

xmon_options = Dict(
    connection_pads=Dict(
        readout = Dict(connector_location = '180')),
)


# Create a new Transmon Cross object with name 'Q1'
TransmonCross(design, 'transmon', options=xmon_options)

# gui.rebuild()  # rebuild the design and plot
# gui.autoscale() #resize GUI to see QComponent

from qiskit_metal.renderers.renderer_ansys_pyaedt.hfss_renderer_eigenmode_aedt import QHFSSEigenmodePyaedt
import pyEPR as epr

em1_aedt = QHFSSEigenmodePyaedt(design, 
                                'Qubit_Sabrina_proj',
                                'Qubit_Sabrina_design',
                                initiate=True)
hfss_aedt = em1_aedt.current_app

def render_design():
    em1_aedt.render_design()

def run_simulation():
    #### Computer's variables
    num_cores=16
    num_gpu=1
    
    #### Setup Name
    setup_name = 'TransmonSetup'
    
    # Change Silicon to Ultracold Silicon
    hfss_aedt.materials['silicon'].permittivity = 11.45
    
    # Add project variables
    add_project_level_vars(em1_aedt)
    
    # Clear setups
    if len(hfss_aedt.setups) != 0:
        hfss_aedt.setups[0].delete()
    
    # Add Setup
    setup_options = dict(
        name= setup_name,
        MinimumFrequency = 3,  # GHz
        NumModes= 1,
        MaxDeltaFreq = 0.1,
        MaximumPasses= 1,
        MinimumPasses= None,
        MinimumConvergedPasses= 2,
        PercentRefinement= 30,
        BasisOrder= None
    )
    
    em1_aedt.add_hfss_em_setup(**setup_options)
    
    hfss_aedt.analyze_setup(name=setup_name, 
                            num_cores=num_cores, 
                            num_gpu=num_gpu)
    
    # Make silicon ultra cold
    hfss_aedt.materials['silicon'].permittivity = 11.45

def add_project_level_vars(renderer):
    """
    Finds all names, inductances, and capacitances of Josephson Junctions rendered into ANSYS.

    Args:
        renderer (QHFSSEigenmodePyaedt)
    """
    # Get all josephson junctions from rendered components table
    geom_table = renderer.path_poly_and_junction_with_valid_comps
    all_jjs = geom_table.loc[geom_table['name'].str.contains('rect_jj')]
    all_jjs = all_jjs.reset_index(drop=True)

    for i, row in all_jjs.iterrows():
        ### Parsing Data ###
        component = str(row['component'])
        name = str(row['name'])
        inductance = row['aedt_hfss_inductance']  # Lj in Henries
        capacitance = row['aedt_hfss_capacitance']  # Cj in Farads

        # Get ANSYS > Model > Sheet corresponding to JJs
        rect_name = 'JJ_rect_Lj_' + component + '_' + name

        # Get ANSYS > Model > Lines corresponding to JJs
        line_name = 'JJ_Lj_' + component + '_' + name + '_'

        ### Appending data ###
        # Add global Lj and Cj variables to ANSYS (for EPR analysis)
        ansys_Lj_name = f'Lj_{i}'
        ansys_Cj_name = f'Cj_{i}'

        renderer.set_variable(ansys_Lj_name, str(inductance * 1E9) + 'nH')
        renderer.set_variable(ansys_Cj_name, str(capacitance * 1E15) + 'fF')

def run_analysis():
    '''
    After finishing simulation, analyze results
    '''
    em1_aedt.run_epr()
    
    # Extract observables of int
    observables = extract_observables(em1_aedt.epr_quantum_analysis)
    return observables



def extract_observables(epra, 
                        cos_trunc: int = 10, 
                        fock_trunc: int = 15):
    '''
    Extract observables from EPR quantum analysis
    
    Args:
        epra (epr.QuantumAnalysis): Raw EPR Quanutm analysis class
        cos_trunc (int, optional): amount of terms to consider in cosine expansion
        fock_trunc (int, optional): size of fock space before truncation
        
    Returns:
        package (dict): simulated observables
    '''
    
    # Extraction of variables
    omegas = epra.get_frequencies()
    chis = epra.get_chis()
    other_data = epra.data
    
    package = dict(qubit_freq_MHz = omegas['0'][0],
                   qubit_anharmonicity_MHz = chis[0][0],
                   other_data = str(other_data))
    
    return package

def run_single_design(): # this will be used as `custom_analysis`

    render_design()
    run_simulation()
    package = run_analysis()
    return package

from metal_library.sweeper import QSweeper
sweeper = QSweeper(design)

sweep_parameters = dict(
  connection_pads = dict(
    readout = dict(
      connector_type = '0', 
      claw_length = ['185um', '195um', '205um', '215um'],
      ground_spacing = ['4um', '5um', '7um', '9um', '10um'], 
      claw_gap = ['5.1um'], 
      claw_width = ['10um', '15um']
      )
    ),
  cross_width = ['30um'],
  cross_length = ['185um', '195um', '205um', '215um', '225um', '235um'],
  cross_gap = ['29um'],
  aedt_hfss_inductance = [9.686E-9, 10.471E-9, 11.268E-9, 12.679-9, 13.816]
)

my_librarian = sweeper.run_single_component_sweep(component_name='transmon',
                                                  parameters = sweep_parameters,
                                                  custom_analysis = run_single_design,
                                                  parameters_slice=slice(1,2),
                                                  save_path="./Tutorial2_QubitOnly.csv")

