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


from __future__ import annotations


__all__ = ['netlist', 'netlist_from_yaml', 'logical_netlist']

# Cell
#nbdev_comment from __future__ import annotations

import os
import re
from functools import partial
from typing import Callable, Dict, Iterable, Optional, Tuple, Union, cast

from sax import FrozenDict
from natsort import natsorted
from .models import models as default_sax_models
from .typing_ import (
    ComplexFloat,
    Instance,
    Instances,
    LogicalNetlist,
    Model,
    ModelFactory,
    Models,
    Netlist,
    Settings,
    is_instance,
    is_model_factory,
    is_netlist,
)
from .utils import (
    clean_string,
    copy_settings,
    get_settings,
    merge_dicts,
    rename_params,
    rename_ports,
    try_float,
)
from yaml import Loader
from yaml import load as load_yaml

# Internal Cell

def _clean_component_names(models: Optional[Models]) -> Models:
    if models is None:
        models = {}
    return {clean_string(comp): model for comp, model in models.items()}


def _clean_instance_names(instances: Optional[Instances]) -> Instances:
    _instances = {}
    if instances is None:
        instances = {}
    for name, inst in instances.items():
        if "," in name:
            raise ValueError(
                f"Instance name '{name}' is invalid. It contains the port separator ','."
            )
        if ":" in name:
            raise ValueError(
                f"Instance name '{name}' is invalid. It contains the port slice symbol ':'."
            )
        name = clean_string(name)
        _instances[name] = inst
    return _instances


def _component_from_callable(f: Callable, models: Models) -> Tuple[str, Models]:
    _reverse_models = {id(model): name for name, model in models.items()}
    if id(f) in _reverse_models:
        component = _reverse_models[id(f)]
    else:
        component = _funcname(f)
        models[component] = f
    return component, models


def _funcname(p: Callable) -> str:
    name = ""
    f: Callable = p
    while isinstance(f, partial):
        name = "{name}p_"
        if f.args:
            try:
                name = f"{name[:-1]}{hash(f.args)}_"
            except TypeError:
                raise TypeError(
                    "when using partials as SAX models, positional arguments of the partial should be hashable."
                )
        f = f.func
    return f"{name}{f.__name__}_{id(f)}"


def _instance_from_callable(f: Callable, models: Models) -> Tuple[Instance, Models]:
    f, settings = _maybe_parse_partial(f)
    component, models = _component_from_callable(f, models)
    instance = Instance(component=component, settings=settings)
    return instance, models


def _instance_from_instance(
    name: str,
    instance: Instance,
    models: Models,
    override_settings: Settings,
    global_settings: Dict[str, ComplexFloat],
    default_models=None,
) -> Tuple[Instance, Models]:
    default_models = default_sax_models if default_models is None else default_models
    component = clean_string(instance["component"])
    if component not in models:
        if component not in default_models:
            raise ValueError(
                f"Error constructing netlist. Component '{component}' not found."
            )
        model = default_models[component]
    else:
        model = models[component]
        if isinstance(model, str):
            if model not in default_models:
                raise ValueError(
                    f"Error constructing netlist. Component '{model}' not found."
                )
            model = default_models[model]
    if not callable(model):  # Model or ModelFactory
        raise ValueError(
            f"Error constructing netlist. Model for component '{component}' is not callable."
        )
    # fmt: off
    _default_settings = get_settings(model)
    _instance_settings = {k: v for k, v in instance.get('settings', {}).items() if k in _default_settings}
    _override_settings = cast(Dict[str, ComplexFloat], override_settings.get(name, {}))
    _override_settings = {k: v for k, v in _override_settings.items() if k in _default_settings}
    _global_settings = {k: v for k, v in global_settings.items() if k in _default_settings}
    settings = merge_dicts(_default_settings, _instance_settings, _override_settings, _global_settings)
    # fmt: on

    instance = Instance(component=component, settings=settings)
    models[instance["component"]] = model
    return instance, models


def _instance_from_string(
    s: str, models: Models, default_models=None
) -> Tuple[Instance, Models]:
    default_models = default_sax_models if default_models is None else default_models
    if s not in models:
        if s not in default_models:
            raise ValueError(f"Error constructing netlist. Component '{s}' not found.")
        models[s] = default_models[s]
    instance = Instance(component=s, settings={})
    return instance, models


def _maybe_parse_partial(p: Callable) -> Tuple[Callable, Dict[str, ComplexFloat]]:
    _settings = {}
    while isinstance(p, partial):
        _settings = merge_dicts(_settings, p.keywords)
        if not p.args:
            p = p.func
        else:
            p = partial(p.func, *p.args)
            break
    return p, _settings


def _model_operations(
    models: Optional[Models],
    ops: Optional[Dict[str, Callable]] = None,
    default_models=None,
) -> Models:
    default_models = default_sax_models if default_models is None else default_models
    if ops is None:
        ops = {}
    if models is None:
        models = {}

    _models = {}
    for component, model in models.items():
        if isinstance(model, str):
            if not model in default_models:
                raise ValueError(f"Could not find model {model}.")
            model = default_models[model]

        _models[component] = model

        if not isinstance(model, dict):
            continue

        if is_netlist(model):
            continue  # TODO: This case should actually be handled...

        model = {**model}

        if "model" not in model:
            raise ValueError(
                "Invalid model dict for '{component}'. Key 'model' not found."
            )

        if isinstance(model["model"], str):
            if not model["model"] in default_models:
                raise ValueError(f"Could not find model {model['model']}.")
            model["model"] = default_models[cast(str, model["model"])]

        for op_name, op_func in ops.items():
            assert isinstance(model, dict)
            if op_name not in model:
                continue
            op_args = model[op_name]
            model["model"] = op_func(model["model"], op_args)

        _models[component] = model["model"]
    return _models


def _split_global_settings(
    settings: Optional[dict], instance_names: Iterable[str]
) -> Tuple[Settings, Dict[str, ComplexFloat]]:
    if settings:
        override_settings = cast(Dict[str, Settings], copy_settings(settings))
        global_settings: Dict[str, ComplexFloat] = {}
        for k in list(override_settings.keys()):
            if k in instance_names:
                continue
            global_settings[k] = cast(ComplexFloat, try_float(override_settings.pop(k)))
    else:
        override_settings: Dict[str, Settings] = {}
        global_settings: Dict[str, ComplexFloat] = {}
    return override_settings, global_settings


def _enumerate_portrange(s):
    if not ":" in s:
        return [s]
    idx1, idx2 = s.split(":")
    s1 = re.sub("[0-9]*", "", idx1)
    idx1 = int(re.sub("[^0-9]*", "", idx1))
    s2 = re.sub("[0-9]*", "", idx2)
    idx2 = int(re.sub("[^0-9]*", "", idx2))
    if s1 != s2 and s2 != "":
        raise ValueError(
            "Cannot enumerate portrange {s}, string portion of port differs."
        )
    return [f"{s1}{i}" for i in range(idx1, idx2)]


def _validate_connections(connections):
    # todo: check if instance names are available in instances
    # todo: check if instance ports are used in output ports
    _ports = set()
    old_connections, connections = connections, {}
    for conn1, conn2 in old_connections.items():
        if conn1.count(",") != 1 or conn2.count(",") != 1:
            raise ValueError(
                "Connections ports should have format '{instance_name},{port}'. "
                f"Got '{conn1}'."
            )
        name1, port1 = conn1.split(",")
        name2, port2 = conn2.split(",")
        ports1 = _enumerate_portrange(port1)
        ports2 = _enumerate_portrange(port2)

        if len(ports1) != len(ports2):
            if len(ports1) == 1:
                ports1 = [ports1[0] for _ in ports2]
            elif len(ports2) == 1:
                ports2 = [ports2[0] for _ in ports1]
            else:
                raise ValueError(
                    f"Cannot enumerate connection {conn1} -> {conn2}, slice lengths on both sides differ."
                )

        for port1, port2 in zip(ports1, ports2):
            name1 = clean_string(name1)
            name2 = clean_string(name2)
            port1 = clean_string(port1)
            port2 = clean_string(port2)
            (name1, port1), (name2, port2) = natsorted([(name1, port1), (name2, port2)])
            conn1 = f"{name1},{port1}"
            conn2 = f"{name2},{port2}"
            if conn1 in _ports:
                raise ValueError(f"duplicate connection port: '{conn1}'")
            if conn2 in _ports:
                raise ValueError(f"duplicate connection port: '{conn2}'")
            connections[conn1] = conn2
            _ports.add(conn1)
            _ports.add(conn2)

    return dict(natsorted([natsorted([k, v]) for k, v in connections.items()]))


def _validate_ports(ports):
    # todo: check if instance names are available in instances
    # todo: check if instance ports are used in connections
    _ports = set()
    old_ports, ports = ports, {}
    for port1, conn2 in old_ports.items():

        if port1.count(",") == 1:
            if conn2.count(",") != 0:
                raise ValueError(
                    "Netlist output port '{conn2}' should not contain any ','."
                )
            port1, conn2 = conn2, port1
        elif conn2.count(",") == 1:
            if port1.count(",") != 0:
                raise ValueError(
                    "Netlist output port '{port1}' should not contain any ','."
                )
        else:
            raise ValueError(
                "Netlist output ports should be mapped onto an instance port "
                "using the format: '{output_port}': '{instance_name},{port}'. "
                f"Got: '{port1}': '{conn2}.'"
            )

        name2, port2 = conn2.split(",")
        ports1 = _enumerate_portrange(port1)
        ports2 = _enumerate_portrange(port2)

        if len(ports1) != len(ports2):
            if len(ports1) == 1:
                ports1 = [ports1[0] for _ in ports2]
            elif len(ports2) == 1:
                ports2 = [ports2[0] for _ in ports1]
            else:
                raise ValueError(
                    f"Cannot enumerate output ports {port1} -> {conn2}, slice lengths on both sides differ."
                )

        for port1, port2 in zip(ports1, ports2):
            port1 = clean_string(port1)  # output_port
            name2 = clean_string(name2)
            port2 = clean_string(port2)
            conn2 = f"{name2},{port2}"
            if port1 in _ports:
                raise ValueError(f"duplicate output port: '{port1}'")
            if conn2 in _ports:
                raise ValueError(f"duplicate instance output port: '{conn2}'")
            ports[port1] = conn2
            _ports.add(port1)
            _ports.add(conn2)
    return dict(natsorted(ports.items()))

# Cell

def netlist(
    *,
    instances: Instances,
    connections: Dict[str, str],
    ports: Dict[str, str],
    models: Optional[Models] = None,
    settings: Optional[Settings] = None,
    default_models=None,
) -> Tuple[Netlist, Models]:
    """Create a `Netlist` and `Models` dictionary"""
    default_models = default_sax_models if default_models is None else default_models
    models = _clean_component_names(models)
    models = _model_operations(
        models,
        ops={
            "rename_params": rename_params,
            "rename_ports": rename_ports,
        },
        default_models=default_models,
    )

    instances = _clean_instance_names(instances)
    override_settings, global_settings = _split_global_settings(settings, instances)

    _instances: Dict[str, Union[Instance, Netlist]] = {}
    for name, instance in instances.items():
        if callable(instance):
            instance, models = _instance_from_callable(instance, models)
        elif isinstance(instance, str):
            instance, models = _instance_from_string(instance, models, default_models)
        if not isinstance(instance, dict):
            raise ValueError(
                f"invalid instance '{name}': expected str, dict or callable."
            )
        if is_instance(instance):
            instance, models = _instance_from_instance(
                name=name,
                instance=cast(Instance, instance),
                models=models,
                override_settings=override_settings,
                global_settings=global_settings,
                default_models=default_models,
            )
        elif is_netlist(instance):
            instance = cast(Netlist, instance)
            instance, models = netlist(
                instances=instance["instances"],
                connections=instance["connections"],
                ports=instance["ports"],
                models=models,
                settings=merge_dicts(global_settings, override_settings),
                default_models=default_models,
            )
        else:
            raise ValueError(
                f"Instance {name} cannot be interpreted as an Instance or a Netlist."
            )
        _instances[name] = instance

    _instances = {k: _instances[k] for k in natsorted(_instances.keys())}
    _connections = _validate_connections(connections)
    _ports = _validate_ports(ports)

    _netlist = Netlist(
        instances=cast(Instances, _instances),
        connections=_connections,
        ports=_ports,
    )
    return _netlist, models

# Cell

def netlist_from_yaml(
    yaml: str,
    *,
    models: Optional[Models] = None,
    settings: Optional[Settings] = None,
    default_models=None,
) -> Tuple[Netlist, Models]:
    """Load a sax `Netlist` from yaml definition and a `Models` dictionary"""

    default_models = default_sax_models if default_models is None else default_models
    ext = None
    directory = None
    yaml_path = os.path.abspath(os.path.expanduser(yaml))
    if os.path.isdir(yaml_path):
        raise IsADirectoryError(
            "Cannot read from yaml path '{yaml_path}'. Path is a directory."
        )
    elif os.path.exists(yaml_path):
        if ext is None:
            _, *ext_list = os.path.basename(yaml_path).split(".")
            ext = f".{'.'.join(ext_list)}"
        if directory is None:
            directory = os.path.dirname(yaml_path)
        yaml = open(yaml_path, "r").read()
    else:
        yaml_path = None

    subnetlists = {}
    if directory is not None and ext is not None:
        subnetlists = {
            re.sub(f"{ext}$", "", os.path.basename(file)): os.path.join(root, file)
            for root, _, files in os.walk(os.path.abspath(directory))
            for file in files
            if file.endswith(ext)
        }

    raw_netlist = load_yaml(yaml, Loader)

    for section in ["instances", "connections", "ports"]:
        if section not in raw_netlist:
            raise ValueError(f"Can not load from yaml: '{section}' not found.")

    raw_instances = raw_netlist["instances"]
    connections = raw_netlist["connections"]
    ports = raw_netlist["ports"]
    override_settings, global_settings = _split_global_settings(settings, raw_instances)

    instances = {}
    for name, instance in raw_instances.items():
        if isinstance(instance, str):
            instance = {"component": instance, "settings": {}}
        elif isinstance(instance, dict):
            if "component" not in instance:
                raise ValueError(
                    f"Can not load from yaml: 'component' not found in instance '{name}'."
                )
        component = instance["component"]
        component = re.sub(r"\.ba$", "", component)  # for compat with pb
        if component in subnetlists:
            _override_settings = cast(Settings, override_settings.get(name, {}))
            _settings = merge_dicts(global_settings, _override_settings)
            instance, models = netlist_from_yaml(
                yaml=subnetlists[component],
                models=models,
                settings=_settings,
                default_models=default_models,
            )
        instances[name] = instance

    _netlist, _models = netlist(
        instances=instances,
        connections=connections,
        ports=ports,
        models=models,
        settings=settings,
        default_models=default_models,
    )
    return _netlist, _models

# Cell

def logical_netlist(
    *,
    instances: Instances,
    connections: Dict[str, str],
    ports: Dict[str, str],
    models: Optional[Models] = None,
    settings: Optional[Settings] = None,
    default_models=None,
) -> Tuple[LogicalNetlist, Settings, Models]:
    """Create a `LogicalNetlist` with separated `Settings` and `Models` dictionary"""

    default_models = default_sax_models if default_models is None else default_models
    _netlist, models = netlist(
        instances=instances,
        connections=connections,
        ports=ports,
        models=models,
        settings=settings,
        default_models=default_models,
    )
    _instances: Dict[str, str] = {}
    _settings: Settings = {}
    _models = models

    for name, instance in _netlist["instances"].items():
        if is_netlist(instance):
            instance = cast(Netlist, instance)
            model, _settings[name], _models = logical_netlist(
                instances=instance["instances"],
                connections=instance["connections"],
                ports=instance["ports"],
                models=_models,
                default_models=default_models,
            )
            model_hash = hex(abs(hash(FrozenDict(model))))[2:]
            component = f"logical_netlist_{model_hash}"
            _instances[name] = component
            _models[component] = model
        elif is_instance(instance):
            instance = cast(Instance, instance)
            component = instance["component"]
            _instance_settings = instance.get("settings", {})
            _instance_model = cast(Model, _models[component])

            if is_model_factory(_instance_model):
                model_factory = cast(ModelFactory, _instance_model)
                _instance_model = model_factory(**_instance_settings)
                instance, _models = _instance_from_callable(_instance_model, _models)
                _instance_settings = instance.get("settings", {})
                component = instance["component"]

            _instances[name] = component
            _model_settings = get_settings(_instance_model)
            _instance_settings = {
                k: try_float(
                    v
                    if (k not in _instance_settings or _instance_settings[k] is None)
                    else _instance_settings[k]
                )
                for k, v in _model_settings.items()
            }
            _settings[name] = cast(Settings, _instance_settings)
        else:
            raise ValueError(f"instance '{name}' is not an Instance or a Netlist.")

    _instances = dict(natsorted(_instances.items()))

    _logical_netlist = LogicalNetlist(
        instances=_instances,
        connections=_netlist["connections"],
        ports=_netlist["ports"],
    )
    return (
        _logical_netlist,
        _settings,
        _models,
    )