# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/08_circuit.ipynb (unless otherwise specified).


from __future__ import annotations


__all__ = ['circuit', 'circuit_from_netlist', 'circuit_from_yaml', 'circuit_from_gdsfactory']

# Cell
#nbdev_comment from __future__ import annotations

from functools import partial
from typing import Dict, Optional, Tuple, Union, cast

from .backends import circuit_backends
from .models import coupler, straight
from .multimode import multimode, singlemode
from .netlist import LogicalNetlist, Netlist, logical_netlist, netlist_from_yaml
from .typing_ import Instances, Model, Models, Netlist, Settings, SType, is_netlist
from .utils import _replace_kwargs, get_settings, merge_dicts, update_settings

try:
    import jax.numpy as jnp
    JAX_AVAILABLE = True
except ImportError:
    import numpy as jnp
    JAX_AVAILABLE = False

# Cell

def circuit(
    *,
    instances: Instances,
    connections: Dict[str, str],
    ports: Dict[str, str],
    models: Optional[Models] = None,
    modes: Optional[Tuple[str, ...]] = None,
    settings: Optional[Settings] = None,
    backend: str = "default",
    default_models=None,
) -> Model:
    # assert valid circuit_backend
    if backend not in circuit_backends:
        raise KeyError(
            f"circuit backend {backend} not found. Allowed circuit backends: "
            f"{', '.join(circuit_backends.keys())}."
        )

    evaluate_circuit = circuit_backends[backend]

    _netlist, _settings, _models = logical_netlist(
        instances=instances,
        connections=connections,
        ports=ports,
        models=models,
        settings=settings,
        default_models=default_models,
    )

    for name in list(_models.keys()):
        if is_netlist(_models[name]):
            netlist_model = cast(LogicalNetlist, _models.pop(name))
            instance_model_names = set(netlist_model["instances"].values())
            instance_models = {k: _models[k] for k in instance_model_names}
            netlist_func = circuit_from_netlist(
                netlist=netlist_model,
                models=instance_models,
                backend=backend,
                modes=modes,
                settings=None,  # settings are already integrated in netlist by now.
                default_models=default_models,
            )
            _models[name] = netlist_func

    if modes is not None:
        maybe_multimode = partial(multimode, modes=modes)
        connections = {
            f"{p1}@{mode}": f"{p2}@{mode}"
            for p1, p2 in _netlist["connections"].items()
            for mode in modes
        }
        ports = {
            f"{p1}@{mode}": f"{p2}@{mode}"
            for p1, p2 in _netlist["ports"].items()
            for mode in modes
        }
    else:
        maybe_multimode = partial(singlemode, mode="te")
        connections = _netlist["connections"]
        ports = _netlist["ports"]

    def _circuit(**settings: Settings) -> SType:
        settings = merge_dicts(_settings, settings)
        global_settings = {}
        for k in list(settings.keys()):
            if k in _netlist["instances"]:
                continue
            global_settings[k] = settings.pop(k)
        if global_settings:
            settings = cast(
                Dict[str, Settings], update_settings(settings, **global_settings)
            )
        instances: Dict[str, SType] = {}
        for name, model_name in _netlist["instances"].items():
            model = cast(Model, _models[model_name])
            instances[name] = cast(
                SType, maybe_multimode(model(**settings.get(name, {})))
            )
        S = evaluate_circuit(instances, connections, ports)
        return S

    settings = {
        name: get_settings(cast(Model, _models[model]))
        for name, model in _netlist["instances"].items()
    }
    settings = merge_dicts(settings, _settings)
    _replace_kwargs(_circuit, **settings)

    return _circuit

# Cell

def circuit_from_netlist(
    netlist: Union[LogicalNetlist, Netlist],
    *,
    models: Optional[Models] = None,
    modes: Optional[Tuple[str, ...]] = None,
    settings: Optional[Settings] = None,
    backend: str = "default",
    default_models=None,
) -> Model:
    """create a circuit model function from a netlist """
    instances = netlist["instances"]
    connections = netlist["connections"]
    ports = netlist["ports"]
    _circuit = circuit(
        instances=instances,
        connections=connections,
        ports=ports,
        models=models,
        modes=modes,
        settings=settings,
        backend=backend,
        default_models=default_models,
    )
    return _circuit

# Cell
def circuit_from_yaml(
    yaml: str,
    *,
    models: Optional[Models] = None,
    modes: Optional[Tuple[str, ...]] = None,
    settings: Optional[Settings] = None,
    backend: str = "default",
    default_models=None,
) -> Model:
    """Load a sax circuit from yaml definition

    Args:
        yaml: the yaml string to load
        models: a dictionary which maps component names to model functions
        modes: the modes of the simulation (if not given, single mode
            operation is assumed).
        settings: override netlist instance settings. Use this setting to set
            global settings like for example the wavelength 'wl'.
        backend: "default" or "klu". How the circuit S-parameters are
            calculated.  "klu" is a CPU-only method which generally speaking is
            much faster for large circuits but cannot be jitted or used for autograd.
    """
    netlist, models = netlist_from_yaml(yaml=yaml, models=models, settings=settings)
    circuit = circuit_from_netlist(
        netlist=netlist,
        models=models,
        modes=modes,
        settings=None,  # settings are already integrated in the netlist by now
        backend=backend,
        default_models=default_models,
    )
    return circuit

# Cell
def circuit_from_gdsfactory(
    component,
    *,
    models: Optional[Models] = None,
    modes: Optional[Tuple[str, ...]] = None,
    settings: Optional[Settings] = None,
    backend: str = "default",
    default_models=None,
) -> Model:
    """Load a sax circuit from a GDSFactory component"""
    circuit = circuit_from_netlist(
        component.get_netlist(),
        models=models,
        modes=modes,
        settings=settings,
        backend=backend,
        default_models=default_models,
    )
    return circuit