"""
networks/cpw.py
---------------
Network builders for CPW (Coplanar Waveguide) resonator topologies.

Each function takes a CPWParams dataclass for device geometry plus
circuit-level parameters (coupling caps, frequency, etc.) and returns
an rf.Network.

Topology overview
-----------------
All networks follow the same chain:

    [Port1] -- Cc1 -- Ctog1 -- [CPW line] -- Ctog2 -- Cc2 -- [Port2 or Open]

The "loaded" variants insert an LC resonator between the port and the
coupling capacitor on each side:

    [Port1] -- cc_port1 -- load1 -- Cc1 -- Ctog1 -- [CPW line] -- ...
"""
from __future__ import annotations
import skrf as rf
from skrf.media import CPW

from simpleLOMs.elements import coupling_capacitor, shunt_capacitor, lc_resonator
from simpleLOMs.params import CPWParams


def _make_cpw_media(freq: rf.Frequency, params: CPWParams) -> CPW:
    """Build a skrf CPW media object from a CPWParams dataclass."""
    return CPW(
        freq,
        w=params.w,
        s=params.s,
        t=params.t,
        h=params.h,
        rho=params.rho,
        ep_r=params.ep_r,
        has_metal_backside=params.has_metal_backside,
        tand=params.tand,
    )


def cpw_resonator_network(
    freq: rf.Frequency,
    d: float,
    Cc1: float,
    Cc2: float,
    Ctog1: float,
    Ctog2: float,
    cpw_params: CPWParams = None,
    Z0: float = 50,
) -> rf.Network:
    """
    Single-port CPW resonator (reflection measurement, port + open termination).

    Topology:
        [Port1] -- Cc1 (series) -- Ctog1 (shunt) -- [CPW line] -- Ctog2 (shunt) -- Cc2 (shunt) -- [Open]

    Parameters
    ----------
    freq : rf.Frequency
    d : float
        Physical resonator length in metres.
    Cc1, Cc2 : float
        Coupling capacitances in Farads.
    Ctog1, Ctog2 : float
        Shunt-to-ground capacitances in Farads.
    cpw_params : CPWParams, optional
        CPW geometry. Uses default CPWParams() if not provided.
    Z0 : float
        Reference impedance in Ohms.

    Returns
    -------
    rf.Network
        1-port reflection network.
    """
    if cpw_params is None:
        cpw_params = CPWParams()

    cpw = _make_cpw_media(freq, cpw_params)
    line   = cpw.line(d=d, unit="m", name="cpw_line")
    cc1    = coupling_capacitor(C=Cc1,  freq=freq, name="cc1",   Z0=Z0)
    cc2    = shunt_capacitor(   C=Cc2,  freq=freq, name="cc2",   Z0=Z0)
    ctog1  = cpw.shunt_capacitor(C=Ctog1, freq=freq, name="ctog1", Z0=Z0)
    ctog2  = cpw.shunt_capacitor(C=Ctog2, freq=freq, name="ctog2", Z0=Z0)
    open_  = rf.Circuit.Open(freq, name="open")
    port1  = rf.Circuit.Port(freq, name="P1", z0=Z0)

    cnx = [
        [(port1, 0), (cc1,   0)],
        [(cc1,   1), (ctog1, 0)],
        [(ctog1, 1), (line,  0)],
        [(line,  1), (ctog2, 0)],
        [(ctog2, 1), (cc2,   0)],
        [(cc2,   1), (open_, 0)],
    ]
    return rf.Circuit(cnx, name="CPW_1port").network


def cpw_resonator_network_2port(
    freq: rf.Frequency,
    d: float,
    Cc1: float,
    Cc2: float,
    Ctog1: float,
    Ctog2: float,
    cpw_params: CPWParams = None,
    Z0: float = 50,
) -> rf.Network:
    """
    Two-port CPW resonator (transmission measurement, port on each side).

    Topology:
        [Port1] -- Cc1 (series) -- Ctog1 (shunt) -- [CPW line] -- Ctog2 (shunt) -- Cc2 (series) -- [Port2]

    Parameters
    ----------
    freq : rf.Frequency
    d : float
        Physical resonator length in metres.
    Cc1, Cc2 : float
        Coupling capacitances in Farads.
    Ctog1, Ctog2 : float
        Shunt-to-ground capacitances in Farads.
    cpw_params : CPWParams, optional
        CPW geometry. Uses default CPWParams() if not provided.
    Z0 : float
        Reference impedance in Ohms.

    Returns
    -------
    rf.Network
        2-port transmission network.
    """
    if cpw_params is None:
        cpw_params = CPWParams()

    cpw = _make_cpw_media(freq, cpw_params)
    line   = cpw.line(d=d, unit="m", name="cpw_line")
    cc1    = coupling_capacitor(C=Cc1,  freq=freq, name="cc1",   Z0=Z0)
    cc2    = coupling_capacitor(C=Cc2,  freq=freq, name="cc2",   Z0=Z0)
    ctog1  = cpw.shunt_capacitor(C=Ctog1, freq=freq, name="ctog1", Z0=Z0)
    ctog2  = cpw.shunt_capacitor(C=Ctog2, freq=freq, name="ctog2", Z0=Z0)
    port1  = rf.Circuit.Port(freq, name="P1", z0=Z0)
    port2  = rf.Circuit.Port(freq, name="P2", z0=Z0)

    cnx = [
        [(port1, 0), (cc1,   0)],
        [(cc1,   1), (ctog1, 0)],
        [(ctog1, 1), (line,  0)],
        [(line,  1), (ctog2, 0)],
        [(ctog2, 1), (cc2,   0)],
        [(cc2,   1), (port2, 0)],
    ]
    return rf.Circuit(cnx, name="CPW_2port").network


def cpw_resonator_loaded_network_2port(
    freq: rf.Frequency,
    d: float,
    Cc1: float,
    Cc2: float,
    Ctog1: float,
    Ctog2: float,
    Lload1: float,
    Cload1: float,
    Lload2: float,
    Cload2: float,
    cpw_params: CPWParams = None,
    Z0: float = 50,
) -> rf.Network:
    """
    Two-port CPW resonator with an LC load on each port side.

    The loads model external resonators (e.g. qubits or readout resonators)
    weakly coupled to each port via small series capacitors.

    Topology:
        [Port1] -- cc_port1 -- load1 (shunt LC) -- Cc1 -- Ctog1 -- [CPW line]
                                                                      -- Ctog2 -- Cc2 -- load2 (shunt LC) -- cc_port2 -- [Port2]

    Parameters
    ----------
    freq : rf.Frequency
    d : float
        Resonator length in metres.
    Cc1, Cc2 : float
        Coupling capacitances in Farads.
    Ctog1, Ctog2 : float
        Shunt-to-ground capacitances in Farads.
    Lload1, Cload1 : float
        Inductance (H) and capacitance (F) of the load on port 1 side.
    Lload2, Cload2 : float
        Inductance (H) and capacitance (F) of the load on port 2 side.
    cpw_params : CPWParams, optional
        CPW geometry. Uses default CPWParams() if not provided.
    Z0 : float
        Reference impedance in Ohms.

    Returns
    -------
    rf.Network
        2-port loaded transmission network.
    """
    if cpw_params is None:
        cpw_params = CPWParams()

    cpw = _make_cpw_media(freq, cpw_params)
    line      = cpw.line(d=d, unit="m", name="cpw_line")
    cc1       = coupling_capacitor(C=Cc1,   freq=freq, name="cc1",      Z0=Z0)
    cc2       = coupling_capacitor(C=Cc2,   freq=freq, name="cc2",      Z0=Z0)
    cc_port1  = coupling_capacitor(C=1e-15, freq=freq, name="cc_port1", Z0=Z0)
    cc_port2  = coupling_capacitor(C=1e-16, freq=freq, name="cc_port2", Z0=Z0)
    ctog1     = cpw.shunt_capacitor(C=Ctog1, freq=freq, name="ctog1", Z0=Z0)
    ctog2     = cpw.shunt_capacitor(C=Ctog2, freq=freq, name="ctog2", Z0=Z0)
    load1     = lc_resonator(L=Lload1, C=Cload1, freq=freq, name="load1", Z0=Z0)
    load2     = lc_resonator(L=Lload2, C=Cload2, freq=freq, name="load2", Z0=Z0)
    port1     = rf.Circuit.Port(freq, name="P1", z0=Z0)
    port2     = rf.Circuit.Port(freq, name="P2", z0=Z0)

    cnx = [
        [(port1,    0), (cc_port1, 0)],
        [(cc_port1, 1), (load1,    0)],
        [(load1,    1), (cc1,      0)],
        [(cc1,      1), (ctog1,    0)],
        [(ctog1,    1), (line,     0)],
        [(line,     1), (ctog2,    0)],
        [(ctog2,    1), (cc2,      0)],
        [(cc2,      1), (load2,    0)],
        [(load2,    1), (cc_port2, 0)],
        [(cc_port2, 1), (port2,    0)],
    ]
    return rf.Circuit(cnx, name="CPW_loaded_2port").network
