#!/usr/bin/env python3
"""
## Assigning names

Names are by default assigned as Abs_0, Abs_1, etc.

For nodes that are called function-style, names can be explicitly assigned:

    node = Gain(name="DemGainz", Gain=2.0)

Or inferred from the Python variable name:

    dem_gainz = Gain(name="DemGainz", Gain=2.0)


## Assigning UUIDs

UUIDs that were randomly generated by the UI are assigned all at once in a call
to `set_uuids`. This is only necessary to roundtrip the UI guids; this call can
in general be deleted harmlessly.

UUIDs that are generated here are generated deterministically, using a hash of
the overall model UUID with the name of the node. This means that that will be
stable across edits, and shouldn't clutter up the code.


## Layout

Layout from the UI is assigned in a block, like UUIDs and node renaming.

Nodes with missing layout info should get some best-effort automatic layout.

We could add some simple layout primitives (e.g. 'row', 'column', 'group') which
the user can wrap their nodes in to control the layout without using
coordinates. If we had alignment utilities in the UI (e.g. select some nodes and
pick 'align in column' or something), these could also be produced on UI->Py.


## Alternatives

Many

"""
import dataclasses

import numbers
import numpy as np
from typing import NamedTuple

from .id import Id
from . import core


def _is_literal(x):
    return isinstance(x, (numbers.Number, np.generic, np.ndarray))


class Bag(dict):
    """
    Wrap dict to allow dot syntax
    """

    def __getattr__(self, name):
        return self[name]


@Id.with_id
class ModelBuilder:
    """
    A collection of blocks.

    Usage:

        m = ModelBuilder("model name")

        block1 = core.Constant(model=m, 3.0)

        # block2 figures out the model from its input, so we
        # don't need to specify it again.
        block2 = core.Gain(in_0 = block1, gain = 2.0)

        # these all do the same thing:
        block5 = core.Adder(inputs=[block1.out_0, block2.out_0])
        block4 = block1.out_0 + block2.out_0
        block3 = block1 + block2

        # name unnnamed blocks from a dictionary
        m.name_from_locals(locals())
    """

    def __init__(self, name="model", allow_overwrites=True):
        self.name = name
        self._name_count = {}
        self.links = {}
        self.nodes = {}
        self.configuration = {}
        self.parameters = {}
        self.groups = {}
        self.submodels = {}
        self._allow_overwrites = allow_overwrites

    def add_parameter(self, key, value, default_value=None, description=""):
        self.parameters[key] = Parameter(
            name=key,
            value=value,
            default_value=default_value,
            description=description,
        )
        return self.parameters[key]

    def update_parameter(self, key, value):
        self.parameters[key].value = value

    def add_group(self, group_name, group_builder):
        group_block = core.Group(name=group_name, input_names=[], output_names=[])
        for node in group_builder.nodes.values():
            if isinstance(node, core.Inport):
                group_block.input_names.append(node.name)
            if isinstance(node, core.Outport):
                group_block.output_names.append(node.name)
        self.groups[group_block] = group_builder
        self.nodes[group_block.name] = group_block
        return group_block

    def add_reference_submodel(self, submodel_name, submodel_builder):
        submodel_block = core.ReferenceSubmodel(
            name=submodel_name, input_names=[], output_names=[]
        )
        for node in submodel_builder.nodes.values():
            if isinstance(node, core.Inport):
                submodel_block.input_names.append(node.name)
            if isinstance(node, core.Outport):
                submodel_block.output_names.append(node.name)
        self.submodels[submodel_block] = submodel_builder
        self.nodes[submodel_block.name] = submodel_block
        return submodel_block

    def add_link(self, src, dst, id=None):
        if isinstance(src, str):
            [node_name, port_name] = src.split(".")
            src = self.nodes[node_name].output_port(port_name)
        if isinstance(dst, str):
            [node_name, port_name] = dst.split(".")
            dst = self.nodes[node_name].input_port(port_name)
        self._ensure_model(src)
        self._ensure_model(dst)

        if isinstance(src, Node):
            src = src.output_port()
        elif _is_literal(src):
            from . import core

            src = core.Constant(self, value=src)
        elif not isinstance(src, OPort):
            raise TypeError(
                f"src should be an output port. Got type {src.__class__.__name__}"
            )

        if isinstance(dst, Node):
            dst = dst.input_port()
        elif not isinstance(dst, IPort):
            raise TypeError(
                f"dst should be an input port. Got type {dst.__class__.__name__}"
            )

        if dst in self.links:
            if not self._allow_overwrites:
                iport_name = f"{dst.node.name}.{dst.name}"
                oport_name = (
                    f"{self.links[dst].src.node.name}.{self.links[dst].src.name}"
                )
                raise TypeError(
                    f"Input port {iport_name} is already connected to output port {oport_name}.\n"
                    "Hint: An input port can only have one link but an output port can have many links."
                )
            else:
                del self.links[dst]

        # pylint: disable=unexpected-keyword-arg
        self.links[dst] = Link(src=src, dst=dst, id=id)
        return self.links[dst]

    def has_link(self, node):
        """
        Return True if a node is connected to another node.
        """
        return any(
            node == link.src.node or node == dst_port.node
            for dst_port, link in self.links.items()
        )

    def new_name(self, base_name):
        "Return a unique node name by appending a number to base_name"
        i = self._name_count.setdefault(base_name, 0)
        self._name_count[base_name] = i + 1
        name = f"{base_name}_{i}"
        while name in self.nodes:
            i += 1
            name = f"{base_name}_{i}"
        return name

    def remove_block(self, block_name: str, remove_links=True):
        """
        Remove a block from the model.
        """
        if block_name not in self.nodes:
            raise ValueError(f"Block {block_name} not found in model")

        if remove_links:
            links_to_remove = set()
            for dst_port in self.links:
                if dst_port.node.name == block_name:
                    links_to_remove.add(dst_port)
                if self.links[dst_port].src.node.name == block_name:
                    links_to_remove.add(dst_port)
            for dst_port in links_to_remove:
                del self.links[dst_port]

        del self.nodes[block_name]

    def remove_link(self, dst):
        """
        Remove a link from the model.
        """
        if not isinstance(dst, IPort):
            raise TypeError(
                f"Argument should be an input port. Got type {dst.__class__.__name__}"
            )
        if dst not in self.links:
            raise ValueError(f"Link for input port {dst} not found in model")
        del self.links[dst]

    def _ensure_model(self, entity):
        if isinstance(entity, Node):
            if entity.model is None:
                self.add_block(entity)
            elif entity.model is not self:
                raise ValueError(f"entity {entity} already has model {entity.model}")
        elif isinstance(entity, IPort):
            self._ensure_model(entity.node)
        elif isinstance(entity, OPort):
            self._ensure_model(entity.node)

    def add_block(self, node):
        if not self._allow_overwrites and node.model is not None:
            raise ValueError(f"node {node} already has model {node.model}")

        if isinstance(node, core.Group):
            if node not in self.groups:
                raise ValueError(
                    "To add a group, call model_builder.add_group(group_name, group_builder)."
                )

        if isinstance(node, core.ReferenceSubmodel):
            if node not in self.submodels:
                raise ValueError(
                    "To add a reference submodel, call "
                    "model_builder.add_reference_submodel(submodel_name, submodel_builder)."
                )

        node.model = self
        if node.name is None:
            node.name = self.new_name(node.schema.name)
        if node.name in self.nodes:
            if not self._allow_overwrites:
                raise KeyError(f"duplicate node name: {node.name}")
            self.remove_block(node.name)

        self.nodes[node.name] = node

    def set_names_from_locals(self, dict):
        """
        Utility to make it easier to assign non-default node names while using
        expressions. You can build models like:

            m = ModelBuilder() a = core.Integrator(..., model=m) b =
            core.DataSource(..., model=m) c = (a + b) / 2.0 return
            m.set_names_from_locals(locals())

        And the names "a", "b", "c" will automatically be assign to the
        corresponding nodes. The node computing (a + b) will still be named
        `Adder_0`.
        """
        nodes = {node.id: node for node in self.nodes.values()}
        for name, node in dict.items():
            if isinstance(node, Node) and node.id in nodes:
                nodes[node.id].name = name
        nodes = {node.name: node for node in self.nodes.values()}
        # TODO: rename existing nodes with name conflicts
        # for now just assert that that didn't happen
        self.nodes.update(nodes)
        return self


class Node:
    def __init__(
        self,
        *args,
        model_builder=None,
        input_names=None,
        output_names=None,
        name=None,
        time_mode=None,
        **params,
    ):
        if args:
            raise TypeError(f"Unexpected positional arguments: {args}")

        self.name = name
        self.pos = None
        self.align = None
        self.model = None
        self.uiprops = {}

        self.input_names = input_names
        self.output_names = output_names
        self.params = params
        self.time_mode = time_mode

        if model_builder is not None:
            model_builder.add_block(self)

    def input_port(self, i=None):
        """`node.input_name(i)` returns the input port at the given index or name.

        `node.input.port()` (with name omitted) requires the node to have exactly
        one input port or one static input port, and returns that port.
        """
        if i is None:
            if len(self.input_names) == 1:
                i = self.input_names[0]
            else:
                i = self.schema.primary_input_port()
                if i is None:
                    raise ValueError(
                        f"Expected a single input port for {self.typename}"
                    )
        elif type(i) is int:
            if i >= len(self.input_names):
                raise ValueError(f"Invalid input port '{i}' for {self.typename}")
            i = self.input_names[i]
        elif i not in self.input_names:
            raise ValueError(f"Invalid input port '{i}' for {self.typename}")
        return IPort(self, i)

    def output_port(self, i=None):
        """`node.output_port(i)` returns the output port at the given index or name.

        `node.output.port()` (with index omitted) requires the node to have exactly
        one output port, and returns that port.
        """
        if i is None:
            if len(self.output_names) == 1:
                i = self.output_names[0]
            else:
                i = self.schema.primary_output_port()
                if i is None:
                    raise ValueError(
                        f"Expected a single output port for {self.typename}"
                    )
        elif type(i) is int:
            if i >= len(self.output_names):
                raise ValueError(f"Invalid output port '{i}' for {self.typename}")
            i = self.output_names[i]
        elif i not in self.output_names:
            raise ValueError(f"Invalid output port '{i}' for {self.typename}")
        return OPort(self, i)

    def update_parameters(self, **kwargs):
        if "input_names" in kwargs:
            self.input_names = kwargs.pop("input_names")
        if "output_names" in kwargs:
            self.output_names = kwargs.pop("output_names")
        self.params.update(kwargs)

    def __getattr__(self, name):
        """Allow accessing ports as e.g. node.out_1"""
        if name in self.input_names:
            return IPort(self, name)
        if name in self.output_names:
            return OPort(self, name)
        error_msg = (
            f'Could not find input/output port "{name}" in {self.name} ({self.__class__.__name__})\n'
            f"Available input_names: {self.input_names}\n"
            f"Available output_names: {self.output_names}\n"
        )
        if self.__class__.__name__ in ["Adder", "Product"] and name.startswith("in_"):
            port_index = int(name[3:])
            op_len = len(self.params["operators"])
            error_msg += (
                f"Hint: {self.name} `operators` only has {op_len} operators, "
                f"but you tried to access in_{port_index}. "
                "Add enough operators to match the number of inputs.\n"
            )
        raise AttributeError(error_msg)

    def __setattr__(self, name, value):
        try:
            params = self.__getattribute__("params")
            if params and name in params:
                self.update_parameters(**{name: value})
            else:
                super().__setattr__(name, value)
        except AttributeError:
            super().__setattr__(name, value)

    def __add__(self, rhs):
        return self.output_port().__add__(rhs)

    def __sub__(self, rhs):
        return self.output_port().__sub__(rhs)

    def __mul__(self, rhs):
        return self.output_port().__mul__(rhs)

    def __truediv__(self, rhs):
        return self.output_port().__truediv__(rhs)

    def __abs__(self):
        return self.output_port().__abs__()


class IPort(NamedTuple):
    """
    Represents an input port of a node.
    """

    node: Node
    name: str

    @property
    def model(self):
        return self.node.model

    @property
    def index(self):
        if self.name not in self.node.input_names:
            raise ValueError(f"Invalid input port '{self.name}' for {self.node}")
        return self.node.input_names.index(self.name)

    def __repr__(self):
        return f'IPort({self.node.id}, "{self.name}")'


class OPort(NamedTuple):
    """
    Represents an output port of a node.

    Output ports behave much like plain values: you can add and subtract them,
    call functions on them, etc., except the functions build up an expression
    (as a graph of nodes) instead of actually evaluating anything.

    You can also explicitly link ports with model.add_link()
    """

    node: Node
    name: str

    @property
    def index(self):
        if self.name not in self.node.output_names:
            raise ValueError(f"Invalid output port '{self.name}' for {self.node}")
        return self.node.output_names.index(self.name)

    def __abs__(self):
        from . import core

        return core.Abs(in_0=self)

    def __add__(self, other):
        if _is_literal(other):
            from . import core

            return core.Offset(in_0=self, offset=other)
        else:
            from . import core

            return core.Adder(in_0=self, in_1=other)

    def __sub__(self, other):
        from . import core

        return core.Adder(in_0=self, in_1=other, operators="+-")

    def __mul__(self, other):
        if _is_literal(other):
            from . import core

            n = core.Gain(self.node.model, gain=other)
            self.node.model.add_link(self, n.in_0)
            return n
        else:
            from . import core

            n = core.Product(self.node.model)
            self.node.model.add_link(self, n.in_0)
            self.node.model.add_link(self, n.in_1)
            return n

    def __truediv__(self, other):
        from . import core

        return core.Product(operators="*/", in_0=self, in_1=other)

    def __lt__(self, other):
        from . import core

        return core.Comparator(operator="<", in_0=self, in_1=other)

    def __le__(self, other):
        from . import core

        return core.Comparator(operator="<=", in_0=self, in_1=other)

    def __gt__(self, other):
        from . import core

        return core.Comparator(operator=">", in_0=self, in_1=other)

    def __ge__(self, other):
        from . import core

        return core.Comparator(operator=">=", in_0=self, in_1=other)

    def __eq__(self, other):
        from . import core

        return core.Comparator(operator="==", in_0=self, in_1=other)

    def __ne__(self, other):
        from . import core

        return core.Comparator(operator="!=", in_0=self, in_1=other)

    def __pow__(self, other):
        from . import core

        return core.Power(in_0=self, Exponent=other)

    @property
    def model(self):
        return self.node.model

    def __repr__(self):
        return f'OPort({self.node.id}, "{self.name}")'


def _node_class_from_schema(schema):
    s = schema

    @Id.with_id
    class C(Node):
        schema = s
        typename = schema.fullname
        input_names = schema.port_names("inputs")
        output_names = schema.port_names("outputs")
        default_params = schema.default_params

        def __init__(self, *args, **kwargs):
            params = C.default_params.copy()
            params.update(kwargs)

            if "input_names" not in kwargs:
                params["input_names"] = C.input_names
            if "output_names" not in kwargs:
                params["output_names"] = C.output_names

            super().__init__(*args, **params)

        @property
        def has_dynamic_input_ports(self):
            return schema.has_dynamic_input_ports

        @property
        def has_dynamic_output_ports(self):
            return schema.has_dynamic_output_ports

    C.__qualname__ = C.__name__ = schema.name  # "Abs"
    C.__module__ = schema.namespace  # "core"
    return C


class Tee(Node):
    """
    Dummy node that represents a tee in a link: it has one input and any
    number of outputs.

    In the editor you can connect a line to the middle of another line with
    alt-click. To do that in a constructed graph, connect `in_0`
    """

    type = "special.Tee"
    schema = None
    input_names = ("in_0",)
    output_names = ("out_0",)


@Id.with_id
@dataclasses.dataclass
class Link:
    src: OPort
    dst: IPort


@Id.with_id
@dataclasses.dataclass
class Parameter:
    name: str
    value: str = None
    default_value: str = None
    description: str = ""
