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

#design.overwrite_enabled = True
#design.chips.main
#design.chips.main.size.size_x = '11mm'
#design.chips.main.size.size_y = '9mm'

# To launch the qiskit metal GUI,use the method MetalGUI.
design = designs.DesignPlanar()
gui = MetalGUI(design)
gui.screenshot()

%metal_heading Connecting QPins with coplanar waveguides (CPWs)

from qiskit_metal.qlibrary.qubits.transmon_pocket import TransmonPocket

# Allow running the same cell here multiple times to overwrite changes
design.overwrite_enabled = True

## Custom options for all the transmons
options = dict(
    # Some options we want to modify from the deafults
    # (see below for defaults)
    pad_width = '425 um', 
    pocket_height = '650um',
    # Adding 4 connectors (see below for defaults)
    connection_pads=dict( 
        a = dict(loc_W=+1,loc_H=+1), 
        b = dict(loc_W=-1,loc_H=+1, pad_height='30um'),
        c = dict(loc_W=+1,loc_H=-1, pad_width='200um'),
        d = dict(loc_W=-1,loc_H=-1, pad_height='50um')
    )
)

## Create 4 transmons

q1 = TransmonPocket(design, 'Q1', options = dict(
    pos_x='+2.55mm', pos_y='+0.0mm', **options))
q2 = TransmonPocket(design, 'Q2', options = dict(
    pos_x='+0.0mm', pos_y='-0.9mm', orientation = '90', **options))
q3 = TransmonPocket(design, 'Q3', options = dict(
    pos_x='-2.55mm', pos_y='+0.0mm', **options))
q4 = TransmonPocket(design, 'Q4', options = dict(
    pos_x='+0.0mm', pos_y='+0.9mm', orientation = '90', **options))

## Rebuild the design
gui.rebuild()
gui.autoscale()

gui.toggle_docks(True)
gui.screenshot()

# Import the basic cpw QComponent from the QLibrary. It is a class called RouteMeander. We can see its default options using RouteMeander.get_template_options(design)
from qiskit_metal.qlibrary.tlines.meandered import RouteMeander
RouteMeander.get_template_options(design)

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', 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='0.13mm'
        ),
        total_length=length,
    fillet = '90um')
    myoptions.update(options)
    myoptions.meander.asymmetry = asymmetry
    myoptions.meander.lead_direction_inverted = 'true' if flip else 'false'
    return RouteMeander(design, component_name, myoptions)


asym = 150
cpw1 = connect('cpw1', 'Q1', 'd', 'Q2', 'c', '6.0 mm', f'+{asym}um')
cpw2 = connect('cpw2', 'Q3', 'c', 'Q2', 'a', '6.1 mm', f'-{asym}um', flip=True)
cpw3 = connect('cpw3', 'Q3', 'a', 'Q4', 'b', '6.0 mm', f'+{asym}um')
cpw4 = connect('cpw4', 'Q1', 'b', 'Q4', 'd', '6.1 mm', f'-{asym}um', flip=True)

gui.rebuild()
gui.autoscale()

gui.toggle_docks(True)
gui.highlight_components(['Q1','Q2','Q3','Q4','cpw1','cpw2','cpw3','cpw4'])
gui.screenshot()

design.components.keys()

# We can access the created CPW from the design too.
design.components.cpw2

# The design can have variables, which can be used in the component options.
design.variables.cpw_width = '10um'
design.variables.cpw_gap = '6um'
gui.rebuild()

# Ex: we can all qubit pads using the variables.
cpw1.options.lead.end_straight = '100um'
cpw2.options.lead.end_straight = '100um'
cpw3.options.lead.end_straight = '100um'
cpw4.options.lead.end_straight = '100um'

# Set variables in the design
design.variables.pad_width = '450 um'
design.variables.cpw_width = '25 um'
design.variables.cpw_gap = '12 um'

# Assign variables to component options
q1.options.pad_width = 'pad_width'
q2.options.pad_width = 'pad_width'
q3.options.pad_width = 'pad_width'
q4.options.pad_width = 'pad_width'

# Rebuild all compoinent and refresh the gui 
gui.rebuild()
gui.autoscale()

gui.screenshot()

%metal_heading Render to GDS

gds = design.renderers.gds
gds.options.path_filename

gds.options.path_filename = '../resources/Fake_Junctions.GDS'

q1.options


gds.options.path_filename = "../resources/Fake_Junctions.GDS"

# Not sure what's going on here, however the GDS file is downloadable from the Qiskit Metal GUI.
design.renderers.gds.export_to_gds("4-Transmon_design.gds")

# Basic and crazy shapes.
# See their source code to see how to get started on a few simple examples.
from qiskit_metal.qlibrary.sample_shapes.n_square_spiral import NSquareSpiral
# print(NSquareSpiral.get_template_options(design))
ops = {
    'n': '10',
    'width': '5um',
    'radius': '100um',
    'gap': '22um',
    'pos_x': '0.65mm',
    'pos_y': '2.2mm',
    'orientation': '0',
    'subtract': 'False'}
NSquareSpiral(design, 'spiral', ops)
NSquareSpiral(design, 'spiral_cut', {**ops, **dict(subtract=True, width='22um', gap='10um')})
gui.rebuild()

# To see source, try ??NSquareSpiral. Go to the actual source file and edit it, copy it, or edit it in the GUI using the Edit Source button.


from qiskit_metal.qlibrary.tlines.straight_path import RouteStraight
# CpwStraightLine.get_template_options(design)
myoptions = Dict(
        pin_inputs=Dict(
            start_pin=Dict(
                component='Q4',
                pin='c'),
            end_pin=Dict(
                component='spiral',
                pin='spiralPin'))
        )


RouteStraight(design, 'cpw_s1', myoptions);
gui.rebuild()

qcomponents = ['spiral', 'cpw_s1']
gui.highlight_components(qcomponents)
gui.zoom_on_components(qcomponents)
gui.screenshot()

# N-Gon version. Orientation is in degrees.
from qiskit_metal.qlibrary.sample_shapes.n_gon import NGon
# display(NGon.get_template_options(design))
ops = {
    'n': '6',
    'radius': '250um',
    'pos_x': '-0.85mm',
    'pos_y': '2.0mm',
    'orientation': '0',
    'subtract': 'False',
    'helper': 'False',
    'chip': 'main',
    'layer': '1'}
NGon(design, 'ngon', ops)
NGon(design, 'ngon_negative', {**ops, **dict(subtract=True, radius='350um')})
gui.rebuild()

gui.zoom_on_components(['ngon_negative'])
gui.screenshot()

# Generate a circle with certain radius.
from qiskit_metal.qlibrary.sample_shapes.circle_raster import CircleRaster
display(CircleRaster.get_template_options(design))

ops = { 'radius': '10um',
        'pos_x': '-1.5mm',
        'pos_y': '2mm',
        'resolution': '16',
        'cap_style': 'round',
        'subtract': 'False',
        'helper': 'False',
        'chip': 'main',
        'layer': '1'}
CircleRaster(design, 'CircleRaster', ops)
gui.rebuild()

gui.zoom_on_components(['CircleRaster'])
gui.screenshot()

# Generate a hollow rectangle. Orientation is in degrees.
from qiskit_metal.qlibrary.sample_shapes.rectangle_hollow import RectangleHollow
display(RectangleHollow.get_template_options(design))

ops = { 'width': '500um',
        'height': '300um',
        'pos_x': '-2.3mm',
        'pos_y': '2mm',
        'orientation': '0',
        'subtract': 'False',
        'helper': 'False',
        'chip': 'main',
        'layer': '1',
        'inner': {  'width': '250um',
                    'height': '100um',
                    'offset_x': '40um',
                    'offset_y': '-20um',
                    'orientation': '15'}}
RectangleHollow(design, 'RectangleHollow', ops)
gui.rebuild()

gui.zoom_on_components(['RectangleHollow'])
gui.screenshot()

gui.autoscale()
gui.screenshot()

# Return the boundry box of the geometry, for example: q1.qgeometry_bounds().
# The function returns a tuple containing (minx, miny, maxx, maxy) bound values for the bounds of the component as a whole.
for name, qcomponent in design.components.items():
    print(f"{name:10s} : {qcomponent.qgeometry_bounds()}")

# We can get all the QGeometry of a QComponent. There are several kinds, such as path and poly. Let us look at all the polygons used to create qubit q1.
q1.qgeometry_table('poly')

q1.qgeometry_table('path')

q1.qgeometry_table('junction')

# Use this code to close the Qiskit Metal GUI. 
gui.main_window.close()

