"""Definitions for describing simulation components."""
from collections import OrderedDict
from enum import Enum
import logging
from typing import Dict, Iterable, List, Optional, Union

import ruamel.yaml as yaml
import yatiml

from ymmsl.identity import Identifier, Reference


_logger = logging.getLogger(__name__)


class Operator(Enum):
    """An operator of a component.

    This is a combination of the Submodel Execution Loop operators,
    and operators for other components such as mappers.
    """

    NONE = 0    #: No operator
    F_INIT = 1  #: Initialisation phase, before start of the SEL
    O_I = 2     #: State observation within the model's main loop
    S = 3       #: State update in the model's main loop
    O_F = 5     #: Observation of final state, after the SEL

    def allows_sending(self) -> bool:
        """Whether ports on this operator can send."""
        return self in {Operator.NONE, Operator.O_I, Operator.O_F}

    def allows_receiving(self) -> bool:
        """Whether ports on this operator can receive."""
        return self in {Operator.NONE, Operator.F_INIT, Operator.S}


class Port:
    """A port on a component.

    Ports are used by component to send or receive messages on. They are
    connected by conduits to enable communication between components.

    Attributes:
        name: The name of the port.
        operator: The MMSL operator in which this port is used.

    """

    def __init__(self, name: Identifier, operator: Operator) -> None:
        """Create a Port.

        Args:
            name: The name of the port.
            operator: The MMSL Operator in which this port is used.
        """
        self.name = name  # type: Identifier
        self.operator = operator  # type: Operator


class Ports:
    """Ports declaration for a component.

    Ports objects compare for equality by value. The names may be
    specified as a list of strings, or separated by spaces in a single
    string. If a particular operator has no associated ports, it may
    be omitted. For example:

    .. code:: yaml

        ports:
          f_init:   # list of names
          - a
          - b
          o_i: c d  # on one line, space-separated
          s: e      # single port
                    # o_f omitted as it has no ports

    Attributes:
        f_init: The ports associated with the F_INIT operator.
        o_i: The ports associated with the O_I operator
        s: The ports associated with the S operator.
        o_f: The ports associated with the O_F operator
    """
    def __init__(
            self, f_init: Union[None, str, List[str]] = None,
            o_i: Union[None, str, List[str]] = None,
            s: Union[None, str, List[str]] = None,
            o_f: Union[None, str, List[str]] = None) -> None:
        """Create a Ports declaration.

        Args:
            f_init: The ports associated with the F_INIT operator.
            o_i: The ports associated with the O_I operator
            s: The ports associated with the S operator.
            o_f: The ports associated with the O_F operator
        """
        def to_list(ports: Union[None, str, List[str]]) -> List[Identifier]:
            if ports is None:
                return list()

            if isinstance(ports, str):
                ports = ports.split(' ')
            return list(map(Identifier, ports))

        self.f_init = to_list(f_init)
        self.o_i = to_list(o_i)
        self.s = to_list(s)
        self.o_f = to_list(o_f)

        all_names = sorted(self.port_names())
        for i in range(len(all_names) - 1):
            if all_names[i] == all_names[i+1]:
                raise RuntimeError(
                        'Invalid ports specification: port "{}" is specified'
                        ' more than once. Port names must be unique within'
                        ' the component.')

    def port_names(self) -> Iterable[Identifier]:
        """Returns an iterable containing the names of all ports."""
        return self.f_init + self.o_i + self.s + self.o_f

    def all_ports(self) -> Iterable[Port]:
        """Returns an iterable containing all ports."""
        return (
                [Port(n, Operator.F_INIT) for n in self.f_init] +
                [Port(name, Operator.O_I) for name in self.o_i] +
                [Port(name, Operator.S) for name in self.s] +
                [Port(name, Operator.O_F) for name in self.o_f])

    def operator(self, port_name: Identifier) -> Operator:
        """Looks up the operator for a given port.

        Args:
            port_name: Name of the port to look up.

        Return:
            The operator for that port.

        Raises:
            KeyError: If no port with this name was found.
        """
        if port_name in self.f_init:
            return Operator.F_INIT

        if port_name in self.o_i:
            return Operator.O_I

        if port_name in self.s:
            return Operator.S

        if port_name in self.o_f:
            return Operator.O_F

        raise KeyError('No port named "{}" was found'.format(port_name))

    _yatiml_defaults = {
            'f_init': [],
            'o_i': [],
            's': [],
            'o_f': []}  # type: Dict[str, Optional[List[str]]]

    @classmethod
    def _yatiml_sweeten(cls, node: yatiml.Node) -> None:
        node.remove_attributes_with_default_values(cls)


class Component:
    """An object declaring a simulation component.

    Simulation components are things like submodels, scale bridges,
    proxies, and any other program that makes up a model. This class
    represents a declaration of a set of instances of a simulation
    component, and it's used to describe which instances are needed to
    perform a certain simulation.

    Attributes:
        name (ymmsl.Reference): The name of this component.
        implementation (ymmsl.Reference): A reference to the
                implementation to use.
        multiplicity (List[int]): The shape of the array of instances
                that execute simultaneously.
        ports (Optional[Ports]): The ports of this component,
                organised by operator. None if not specified.

    """

    def __init__(self, name: str, implementation: Optional[str] = None,
                 multiplicity: Union[None, int, List[int]] = None,
                 ports: Optional[Ports] = None) -> None:
        """Create a Component.

        Args:
            name: The name of the component; must be a valid
                    Reference.
            implementation: The name of the implementation; must be a
                    valid Reference.
            multiplicity: A list of ints describing the shape of the
                    set of instances.
            ports: The ports used by this component to communicate,
                    organised by operator.

        """
        self.name = Reference(name)
        if implementation is None:
            self.implementation = None      # type: Optional[Reference]
        else:
            self.implementation = Reference(implementation)
            for part in self.implementation:
                if isinstance(part, int):
                    raise ValueError('Component implementation {} contains a'
                                     ' subscript, which is not'
                                     ' allowed.'.format(self.name))

        if multiplicity is None:
            self.multiplicity = list()
        elif isinstance(multiplicity, int):
            self.multiplicity = [multiplicity]
        else:
            self.multiplicity = multiplicity

        self.ports = ports

    def __str__(self) -> str:
        """Returns a string representation of the object."""
        result = str(self.name)
        for dim in self.multiplicity:
            result += '[0:{}]'.format(dim)
        return result

    def __repr__(self) -> str:
        """Returns a string representation of the object."""
        return 'Component({})'.format(self.name)

    def instances(self) -> List[Reference]:
        """Creates a list of instances needed.

        Returns:
            A list with one Reference for each instance of this
            component.
        """
        def increment(index: List[int], dims: List[int]) -> None:
            # assumes index and dims are the same length > 0
            # modifies index argument
            i = len(index) - 1
            index[i] += 1
            while index[i] == dims[i]:
                index[i] = 0
                i -= 1
                if i == -1:
                    break
                index[i] += 1

        def generate_indices(multiplicity: List[int]) -> List[List[int]]:
            # n-dimensional counter
            indices = list()    # type: List[List[int]]

            index = [0] * len(multiplicity)
            indices.append(index.copy())
            increment(index, multiplicity)
            while sum(index) > 0:
                indices.append(index.copy())
                increment(index, multiplicity)
            return indices

        result = list()     # type: List[Reference]

        if len(self.multiplicity) == 0:
            result.append(self.name)
        else:
            for index in generate_indices(self.multiplicity):
                result.append(self.name + index)

        return result

    def _yatiml_attributes(self) -> OrderedDict:
        """Reorder attributes for saving to YAML."""
        return OrderedDict([
            ('name', self.name),
            ('ports', self.ports),
            ('multiplicity', self.multiplicity),
            ('implementation', self.implementation)])

    @classmethod
    def _yatiml_recognize(cls, node: yatiml.UnknownNode) -> None:
        node.require_attribute('name', str)

    @classmethod
    def _yatiml_savorize(cls, node: yatiml.Node) -> None:
        if node.has_attribute('multiplicity'):
            if node.has_attribute_type('multiplicity', int):
                attr = node.get_attribute('multiplicity')
                start_mark = attr.yaml_node.start_mark
                end_mark = attr.yaml_node.end_mark
                new_seq = yaml.nodes.SequenceNode(
                        'tag:yaml.org,2002:seq', [attr.yaml_node], start_mark,
                        end_mark)
                node.set_attribute('multiplicity', new_seq)

    @classmethod
    def _yatiml_sweeten(cls, node: yatiml.Node) -> None:
        multiplicity = node.get_attribute('multiplicity')
        items = multiplicity.seq_items()
        if len(items) == 0:
            node.remove_attribute('multiplicity')
        elif len(items) == 1:
            node.set_attribute('multiplicity', items[0].get_value())

        ports_node = node.get_attribute('ports').yaml_node
        if len(ports_node.value) == 0:
            node.remove_attribute('ports')
