# This code is part of KQCircuits
# Copyright (C) 2021 IQM Finland Oy
#
# This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public
# License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied
# warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along with this program. If not, see
# https://www.gnu.org/licenses/gpl-3.0.html.
#
# The software distribution should follow IQM trademark policy for open-source software
# (meetiqm.com/iqm-open-source-trademark-policy). IQM welcomes contributions to the code.
# Please see our contribution agreements for individuals (meetiqm.com/iqm-individual-contributor-license-agreement)
# and organizations (meetiqm.com/iqm-organization-contributor-license-agreement).


from kqcircuits.simulations.simulation import Simulation
from kqcircuits.pya_resolver import pya
from kqcircuits.util.parameters import Param, pdt, add_parameters_from
from kqcircuits.chips.xmons_direct_coupling import XMonsDirectCoupling
from kqcircuits.simulations.port import InternalPort
from kqcircuits.elements.fluxlines.fluxline import Fluxline
from kqcircuits.junctions.junction import Junction


@add_parameters_from(Junction, junction_type="Sim")
@add_parameters_from(Fluxline, fluxline_type="none")
@add_parameters_from(XMonsDirectCoupling, "arm_width_a", "rr_cpl_width")
class XMonsDirectCouplingSim(Simulation):

    # Re-define some XMonsDirectCoupling parameters with defaults changed:
    qubit_spacing = Param(pdt.TypeDouble, "Qubit spacing", 3, unit="μm")
    arm_width_b = Param(pdt.TypeDouble, "Qubit 2 arm width", 66, unit="μm")
    waveguide_length = Param(
        pdt.TypeDouble, "Length of waveguide stubs or distance between couplers and waveguide turning point", 100
    )
    cpl_width = Param(pdt.TypeDouble, "Qubit RR coupler width", 24, unit="μm")
    junction_inductances = Param(
        pdt.TypeList, "Qubit junction inductances", [13.5e-9, 13.5e-9, 13.5e-9], unit="[H, H, H]"
    )
    junction_capacitances = Param(
        pdt.TypeList, "Qubit junction capacitances", [0.1e-15, 0.1e-15, 0.1e-15], unit="[F, F, F]"
    )

    produce_qubits = XMonsDirectCoupling.produce_qubits

    def build(self):

        # parameter used by produce_qubits
        self.rr_cpl_width = (self.cpl_width, self.cpl_width, self.cpl_width)

        self.produce_qubits()

        # driveline 1
        self.produce_waveguide_to_port(
            port_nr=1,
            location=self.refpoints["QB1_port_drive"],
            term1=10,
            towards=self.refpoints["QB1_port_drive"] + pya.DVector(-1, 1),
        )
        # driveline 2
        self.produce_waveguide_to_port(
            port_nr=2,
            location=self.refpoints["QB2_port_drive"],
            term1=10,
            towards=self.refpoints["QB2_port_drive"] + pya.DVector(0, -1),
        )
        # driveline 3
        self.produce_waveguide_to_port(
            port_nr=3,
            location=self.refpoints["QB3_port_drive"],
            term1=10,
            towards=self.refpoints["QB3_port_drive"] + pya.DVector(1, 1),
        )

        # readout 1
        self.produce_waveguide_to_port(
            port_nr=4,
            a=5,
            b=20,
            location=self.refpoints["QB1_port_cplr1"],
            towards=self.refpoints["QB1_port_cplr1"] + pya.DVector(0, 1),
            use_internal_ports=True,
        )
        # readout 2
        self.produce_waveguide_to_port(
            port_nr=5,
            a=5,
            b=20,
            location=self.refpoints["QB2_port_cplr1"],
            towards=self.refpoints["QB2_port_cplr1"] + pya.DVector(0, 1),
            use_internal_ports=True,
        )
        # readout 3
        self.produce_waveguide_to_port(
            port_nr=6,
            a=5,
            b=20,
            location=self.refpoints["QB3_port_cplr1"],
            towards=self.refpoints["QB3_port_cplr1"] + pya.DVector(0, 1),
            use_internal_ports=True,
        )

        # qubits
        self.ports.extend(
            [
                InternalPort(
                    7,
                    *self.etched_line(self.refpoints["QB1_port_squid_a"], self.refpoints["QB1_port_squid_b"]),
                    inductance=self.junction_inductances[0],
                    capacitance=self.junction_capacitances[0],
                    junction=True,
                ),
                InternalPort(
                    8,
                    *self.etched_line(self.refpoints["QB2_port_squid_a"], self.refpoints["QB2_port_squid_b"]),
                    inductance=self.junction_inductances[1],
                    capacitance=self.junction_capacitances[1],
                    junction=True,
                ),
                InternalPort(
                    9,
                    *self.etched_line(self.refpoints["QB3_port_squid_a"], self.refpoints["QB3_port_squid_b"]),
                    inductance=self.junction_inductances[2],
                    capacitance=self.junction_capacitances[2],
                    junction=True,
                ),
            ]
        )
