#!/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 numbers
import numpy as np
from typing import NamedTuple, Optional

from .id import Id


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"):
        self.name = name
        self._name_count = {}
        self.links = {}
        self.nodes = {}
        self.configuration = {}
        self.parameters = {}
        self.groups = {}
        self.submodels = {}
        self.inports = []
        self.outports = []

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

    def add_group(self, node, group):
        node.input_names = group.inports
        node.output_names = group.outports
        self.groups[node] = group

    def add_reference_submodel(self, node, submodel):
        node.input_names = submodel.inports
        node.output_names = submodel.outports
        self.submodels[node] = submodel

    def add_link(self, src, dst):
        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"Expected src is OPort for add_link; got {src}")

        if isinstance(dst, Node):
            dst = dst.input_port()
        elif not isinstance(dst, IPort):
            raise TypeError(f"Expected dst is IPort for add_link; got {dst}")

        if dst in self.links:
            raise TypeError(f"Link {dst} is already connected to {self.links[dst]}")
        self.links[dst] = Link(src=src, dst=dst)
        return self.links[dst]

    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
        return f"{base_name}_{i}"

    def _ensure_model(self, entity):
        if isinstance(entity, Node):
            if entity.model is None:
                self._add(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(self, node):
        assert node.model is None
        node.model = self
        if node.name is None:
            node.name = self.new_name(node.schema.name)
        if node.name in self.nodes:
            raise KeyError(f"duplicate node name: {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
        assert len(nodes) == len(self.nodes)
        self.nodes = nodes
        return self


class Node:
    def __init__(self, model_builder=None, input_names=None, output_names=None, name=None, time_mode=None, **params):
        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(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 __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)
        raise AttributeError(
            f"Could not find {name} in {self.name} ({self.__class__.__name__})\n"
            f"input_names: {self.input_names}\n"
            f"output_names: {self.output_names}\n"
        )

    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__()


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

    node: Node
    name: str
    uuid: Optional[str] = None

    @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}")'


@Id.with_id
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
    uuid: Optional[str] = None

    def __hash__(self):
        return self.uuid.__hash__()

    @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
class Link(NamedTuple):
    src: OPort
    dst: IPort
    pass
