#  Copyright 2020 Modelling, Simulation and Design Lab (MSDL) at
#  McGill University and the University of Antwerp (http://msdl.cs.mcgill.ca/)
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""Set of PythonPDEVS building blocks for transforming data.

The classes in this file concern themselves with message
transformations throughout a system and can come in useful in
numerous situations.
"""

from pypdevs.DEVS import AtomicDEVS
from pypdevs.infinity import INFINITY


class Transformer(AtomicDEVS):
    """Transforms data w.r.t. a function.

    Note:
        The transformation itself happens upon arrival of a
        new item. The more expensive the transformation function, the
        more expensive the usage of this block. Keep this in mind when
        modeling for real-time simulations.

    .. image:: ../img/transform/transformer.svg
		:align: center

    Args:
          name (str):       The name of the block.
          T (def):          Function that defines the transformation.
                            It takes 3 arguments: `i` (the item), `e`
                            (the elapsed time) and `t` (the total time).
                            And it will return the transformed item.
                            Defaults to the item itself.

    State:
        time (numeric):     The total elapsed time.
        transformed (any):  The transformed item.
        original (any):     The original item (before transformation).

    Input Ports:
        input (any):        The item to transform.

    Output Ports:
        transformed (any):  The transformed item.
        original (any):     The original item (before transformation).
    """
    def __init__(self, name: str, T = lambda i, e, t: i):
        super(Transformer, self).__init__(name)
        self.state = {
            "time": 0.0,
            "transformed": None,
            "original": None
        }
        self.T = T

        self.input = self.addInPort("input")
        self.transformed = self.addOutPort("transformed")
        self.original = self.addOutPort("original")

    def extTransition(self, inputs):
        self.state["time"] += self.elapsed
        if self.input in inputs:
            value = inputs[self.input]
            self.state["original"] = value
            self.state["transformed"] = self.T(value, self.elapsed, self.state["time"])
        return self.state

    def intTransition(self):
        self.state["transformed"] = self.state["original"] = None
        return self.state

    def timeAdvance(self):
        if self.state["transformed"] is not None or self.state["original"] is not None:
            return 0
        return INFINITY

    def outputFnc(self):
        res = {}
        if self.state["transformed"] is not None:
            res[self.transformed] = self.state["transformed"]
        if self.state["original"] is not None:
            res[self.original] = self.state["original"]
        return res


# TODO: Interpolation?
class LookupTable(Transformer):
    """Finds the value in a lookup table.

    .. image:: ../img/transform/lookup-table.svg
		:align: center

    Args:
        name (str):         The name of the block.
        table (dict):       The lookup table. Each input is mapped on a
                            key in this table, which is outputted.
        default (any):      The default value if no mapping exists.
                            Defaults to ``None``.

    State:
        time (numeric):     The total elapsed time.
        transformed (any):  The transformed item.
        original (any):     The original item (before transformation).

    Input Ports:
        input (any):        The item to look up.

    Output Ports:
        transformed (any):  The looked up item.
        original (any):     The original item (before lookup).
    """
    def __init__(self, name: str, table: dict, default=None):
        super(LookupTable, self).__init__(name)
        self.table = table
        self.default = default
        self.T = lambda i, _, __: self.table.get(i, self.default)


from pypdevsbbl.extra.packaging import Package
class Pack(AtomicDEVS):
    """Collects all input in a generic package.

    A group of items is captured on this block's input. Based on a
    predefined package (see :mod:`pypdevsbbl.extra.packaging` submodule),
    the elapsed time and the simulation time, there can be deduced
    how the items can be combined.

    As soon as a message arrives on the `release` input, the package
    is released on the `out` port and the state is reset. Additionally,
    if the :meth:`pypdevsbbl.extra.packaging.Package.finished` function of
    evaluates ``True`` after packing, the package will be released
    automatically.

    .. image:: ../img/transform/pack.svg
		:align: center

    Args:
        name (str):         The name of the block.
        package:            The package class in which to combine all input.
                            Defaults to a generic :class:`pypdevsbbl.extra.packaging.Package`.
        *args:              Additional list of arguments to be passed to the
                            Package class upon construction. This is empty
                            by default.

    State:
        package (Package):  See :class:`pypdevsbbl.extra.packaging.Package`.
        release (bool):     Whether or not a release was requested.
        time (numeric):     The total simulation time.

    Input Ports:
        input (any):        The object to add to the package.
        release (any):      When the package needs to be released.

    Output Ports:
        output (Package):   A package.
    """
    def __init__(self, name: str, package = Package, *args):
        super(Pack, self).__init__(name)
        self.cls = package
        self.args = args
        self.state = {
            "package": self.cls(*self.args),
            "release": False,
            "time": 0.0
        }

        self.input = self.addInPort("input")
        self.release = self.addInPort("release")
        self.output = self.addOutPort("output")

    def timeAdvance(self):
        if self.state["release"]:
            return 0
        return INFINITY

    def intTransition(self):
        self.state["release"] = False
        self.state["package"] = self.cls(*self.args)
        return self.state

    def outputFnc(self):
        if self.state["release"]:
            return { self.output: self.state["package"] }
        return {}

    def extTransition(self, inputs):
        self.state["time"] += self.elapsed
        if self.release in inputs:
            self.state["release"] = True
        if self.input in inputs:
            self.state["package"].pack(inputs[self.input], self.elapsed, self.state["time"])

            # Automatic release:
            self.state["release"] = self.state["package"].finished()
        return self.state


class Unpack(AtomicDEVS):
    """Releases the contents of a package.

    When a :class:`pypdevsbbl.extra.packaging.Package` arrives on the
    `input`, it is unpacked and the list of contents is afterwards
    outputted on the `out` port.

    .. image:: ../img/transform/unpack.svg
		:align: center

    Args:
        name (str):         The name of the block.

    State:
        contents (list):    The list of contents of the package.

    Input Ports:
        input (Package):    The package to unpack.

    Output Ports:
        output (list):      All contents of the package.
    """
    def __init__(self, name: str):
        super(Unpack, self).__init__(name)
        self.state = {
            "contents": []
        }

        self.input = self.addInPort("input")
        self.output = self.addOutPort("output")

    def timeAdvance(self):
        if len(self.state["contents"]) > 0:
            return 0
        return INFINITY

    def intTransition(self):
        self.state["contents"].clear()
        return self.state

    def outputFnc(self):
        if len(self.state["contents"]) > 0:
            return { self.output: self.state["contents"][:] }
        return {}

    def extTransition(self, inputs):
        if self.input in inputs:
            self.state["contents"] = inputs[self.input].unpack()
        return self.state
