#  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 doing math.

The core of this module is the :class:`Equation` block,
which allows a user to execute a function on a set of
inputs.

Additionally, some few blocks with a variable set of inputs
(like the sum and the product) will provide the remaining
functionality.
"""

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


class _SP(AtomicDEVS):
    """Superclass for the :class:`Adder` and :class:`Multiplier` blocks.

    Warning:
        This class should not be used as a block by itself.

    Args:
        name (str):         The name of the block.
        n (int):            The amount of input ports the block must have.
                            These ports are accessible from the :attr:`inputs`
                            list, using a 0-based index.

    State:
        result (numeric):   The computed result of the block.
        input-n (numeric):  The value of the `n` th input port.

    Input Ports:
        inputs (list):      The list of all the input ports.

    Output Ports:
        result (numeric):   The result of the computation.
    """
    def __init__(self, name, n):
        super(_SP, self).__init__(name)
        self.inputs = []
        self.state = {
            "result": None
        }
        for i in range(n):
            self.inputs.append(self.addInPort("input-%i" % i))
            self.state["input-%i" % i] = None
        self.result = self.addOutPort("result")

    def operation(self):
        """The operation that needs to be computed upon arrival of any value."""
        raise NotImplementedError()

    def timeAdvance(self):
        if self.state["result"] is not None:
            return 0.0
        return INFINITY

    def outputFnc(self):
        if self.state["result"] is not None:
            return { self.result: self.state["result"] }
        return {}

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

    def extTransition(self, inputs):
        for port in inputs:
            self.state[port.name] = inputs[port]
            if len([1 for x in self.state if x != "result" and self.state[x] is None]) == 0:
                self.state["result"] = self.operation()
        return self.state


class Adder(_SP):
    """Computes the sum of a set of inputs.

    .. image:: ../img/math/adder.svg
		:align: center

    Args:
        name (str):         The name of the block.
        n (int):            The amount of input ports the block must have.
                            These ports are accessible from the :attr:`inputs`
                            list, using a 0-based index.

    State:
        result (numeric):   The computed result of the block.
        input-n (numeric):  The value of the `n` th input port.

    Input Ports:
        inputs (list):      The list of all the input ports.

    Output Ports:
        result (numeric):   The result of the computation.
    """
    def operation(self):
        """The actual sum operation."""
        res = 0
        for port in self.state:
            if port == "result" or self.state[port] is None: continue
            res += self.state[port]
        return res


class Multiplier(_SP):
    """Computes the product of a set of inputs.

    .. image:: ../img/math/multiplier.svg
		:align: center

    Args:
        name (str):         The name of the block.
        n (int):            The amount of input ports the block must have.
                            These ports are accessible from the :attr:`inputs`
                            list, using a 0-based index.

    State:
        result (numeric):   The computed result of the block.
        input-n (numeric):  The value of the `n` th input port.

    Input Ports:
        inputs (list):      The list of all the input ports.

    Output Ports:
        result (numeric):   The result of the computation.
    """
    def operation(self):
        """The actual product operation."""
        res = 1
        for port in self.state:
            if port == "result" or self.state[port] is None: continue
            res *= self.state[port]
        return res


class Equation(AtomicDEVS):
    """Execute an equation on a set of inputs.

    The input ports of this block are determined by the arguments of
    the passed :attr:`function`. Whenever a change happens on these
    ports and both have values, their result is outputted on the
    `result` port.

    Warning:
        - The :attr:`function` cannot have a ``result`` parameter to
          prevent ambiguity with the output port.
        - The :attr:`function` **must not** change the values of the
          inputs.
        - This block is optimized for numerical equations, yet it is
          (because of *Python*) perfectly possible to use some other
          attributes.

    See Also:
        :mod:`pypdevsbbl.extra.equations` for a set of predefined equations which
        can be used as a function without too much thought.

    .. image:: ../img/math/equation.svg
		:align: center

    Args:
        name (str):     The name of the block.
        function (def): The function that needs to be executed. Note that
                        defaults can be used, but the ``*args`` and
                        ``**kwargs`` information will be ignored.

    :Input Ports:       Depends on the :attr:`function` that was passed.
                        Say the :attr:`function` has two arguments: ``a``
                        and ``b``, the input ports are also `a` and `b`.

    Output Ports:
        result (any):   The output of the provided equation.
    """
    def __init__(self, name, function):
        super(Equation, self).__init__(name)
        self.function = function
        self.signature = inspect.signature(self.function)
        self.state = {
            "result": None
        }
        self._ports = []
        for name in self.signature.parameters:
            assert name != "result"
            param = self.signature.parameters[name]
            self._ports.append(name)
            setattr(self, name, self.addInPort(name))
            self.state[name] = param.default
        self.result = self.addOutPort("result")

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

    def outputFnc(self):
        if self.state["result"] is not None:
            return { self.result: self.state["result"] }
        return {}

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

    def extTransition(self, inputs):
        for arg in inputs:
            self.state[arg.name] = inputs[arg]
            if len([0 for p in self._ports if self.state[p] is inspect.Parameter.empty]) == 0:
                self.state["result"] = self.function(**{p: self.state[p] for p in self._ports})
        return self.state


class _DI(AtomicDEVS):
    """Superclass of the :class:`Differentiator` and :class:`Integrator` blocks.

    Warning:
        This class should not be used as a block by itself.

    Args:
        name (str):         The name of the block.
        ic (numeric):       The initial condition of the block. Defaults to 0.

    State:
        prev (numeric):     The previous value. Defaults to the initial condition.
        v (numeric):        The current value.

    Input Ports:
        input (numeric):    The new value.

    Output Ports:
        result (numeric):   The differential.
    """
    def __init__(self, name, ic = 0):
        super(_DI, self).__init__(name)
        self.state = {
            "prev": ic,
            "v": None
        }
        self.input = self.addInPort("input")
        self.result = self.addOutPort("result")

    def timeAdvance(self):
        if self.state["v"] is not None:
            return 0.0
        return INFINITY

    def outputFnc(self):
        if self.state["v"] is not None:
            return { self.result: self.state["v"] }
        return {}

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

    def extTransition(self, inputs):
        raise NotImplementedError()


class Differentiator(_DI):
    """Differentiates the input over time.

    This block will make use of the backwards difference:

    :math:`\\dfrac{d}{dt}f(t) = \\lim_{\\Delta t\\to 0}\\dfrac{f(t) - f(t - \\Delta t)}{\\Delta t} \
            = \\dfrac{x_i - x_{i-1}}{\\Delta t}`


    .. image:: ../img/math/differentiator.svg
		:align: center

    Args:
        name (str):         The name of the block.
        ic (numeric):       The initial condition of the block. Defaults to 0.

    State:
        prev (numeric):     The previous value of the differential. Defaults to the initial condition.
        v (numeric):        The current differential.

    Input Ports:
        input (numeric):    The new value.

    Output Ports:
        result (numeric):   The differential.
    """
    def extTransition(self, inputs):
        if self.input in inputs:
            curr = inputs[self.input]
            if self.elapsed == 0:
                self.state["v"] = self.state["prev"]
            else:
                self.state["v"] = (curr - self.state["prev"]) / self.elapsed
            self.state["prev"] = curr
        return self.state


class Integrator(_DI):
    """Integrates the input over time.

    :math:`v(t) = {\\displaystyle \\int} f(\\tau)\\ d\\tau = ic + {\\displaystyle \\int_0^t} f(\\tau)\\ d\\tau = \
            v(t - \\Delta t) + {\\displaystyle \\int_{t - \\Delta t}^t} f(\\tau)\\ d\\tau`

    .. image:: ../img/math/integrator.svg
		:align: center

    Args:
        name (str):         The name of the block.
        ic (numeric):       The initial condition of the block. Defaults to 0.
        method (def):       A function to determine the computation of the integral.
                            It's preferrably one of :meth:`INTEGRAL_LOWER`, :meth:`INTEGRAL_UPPER` or
                            :meth:`INTEGRAL_TRAPEZOID`. Defaults to :meth:`INTEGRAL_TRAPEZOID`.

    State:
        prev (numeric):     The previous value of the integral. Defaults to the initial condition.
        v (numeric):        The current integral.
        fprev (numeric):    The previous value of the function. When ``None``, the :attr:`method` is
                            ignored.

    Input Ports:
        input (numeric):    The new value.

    Output Ports:
        result (numeric):   The differential.
    """
    def INTEGRAL_LOWER(self, curr):
        """Computes the integral from the lower-bound.

        :math:`{\\displaystyle \\int_{t - \\Delta t}^t} f(\\tau)\\ d\\tau = \\Delta t\\cdot f(t - \\Delta t)`

        Args:
            curr (numeric): Current input value.
        """
        return self.elapsed * self.state["fprev"]

    def INTEGRAL_UPPER(self, curr):
        """Computes the integral from the upper-bound.

        :math:`{\\displaystyle \\int_{t - \\Delta t}^t} f(\\tau)\\ d\\tau = \\Delta t\\cdot f(t)`

        Args:
            curr (numeric): Current input value.
        """
        return self.elapsed * curr

    def INTEGRAL_TRAPEZOID(self, curr):
        """Computes the integral with the trapezoidal rule.

        :math:`{\\displaystyle \\int_{t - \\Delta t}^t} f(\\tau)\\ d\\tau = \
            \\dfrac{\\Delta t\\cdot\\left(f(t - \\Delta t) + f(t)\\right)}{2}`

        Args:
            curr (numeric): Current input value.
        """
        return self.elapsed * (self.state["fprev"] + curr) / 2

    def __init__(self, name, ic = 0, method=INTEGRAL_TRAPEZOID):
        super(Integrator, self).__init__(name, ic)
        self.method = method
        self.state["fprev"] = None

    def extTransition(self, inputs):
        if self.input in inputs:
            curr = float(inputs[self.input])
            if self.state["fprev"] is None:
                self.state["v"] = self.state["prev"] # initial condition
            else:
                self.state["v"] = self.state["prev"] + self.method(self, curr)
            self.state["fprev"] = curr
            self.state["prev"] = self.state["v"]
        return self.state


from pypdevsbbl.extra.rngstreams import createStream, U01
class Random(AtomicDEVS):
    """:class:`AtomicDEVS` that generates a random number every time a message arrives.

    It makes use of the RngStreams random number generator i.i.d. (0, 1).
    Using the ``dist`` argument, one can simply change the distribution over which
    numbers need to be generated. The :mod:`pypdevsbbl.extra.rngstreams.distributions`
    module provides an extensive implementation of a lot of distributions.

    See Also:
        - :class:`pypdevsbbl.generic.generators.RandomNumberGenerator`
        - :class:`pypdevsbbl.generic.generators.RandomDelayGenerator`
        - :mod:`pypdevsbbl.extra.rngstreams`
        - :mod:`pypdevsbbl.extra.rngstreams.distributions`

    .. image:: ../img/math/random.svg
        :align: center

    Args:
        name (str):			The name of the block and the RNG.
        seed (numeric):		The seed for the generator. When ``None``, no seed will be
                            set. Defaults to ``None``.
        dist (def):			A stateless function that implements an inversive cumulative
                            distribution function. It takes at least a single argument,
                            which falls in (0, 1).
                            Defaults to a uniform distribution over (0, 1).
        args (tuple):		Optional additional arguments to add to the distribution.

    State:
        num (numeric):		The random number that's being/needs to be outputted.
        item (any):         The item that's passing through.

    Input Ports:
        input (any):        A new input, indicative of an random number request.

    Output Ports:
        random (numeric):   The random number.
        value (any):	    Outputs the input once more.
    """
    def __init__(self, name: str, seed=None, dist=lambda x: x, args=()):
        super(Random, self).__init__(name)
        self.dist = dist
        self.args = args
        self.gen = createStream(name, seed)
        self.state = {
            "num": self.next(),
            "item": None
        }
        self.input = self.addInPort("input")
        self.random = self.addOutPort("random")
        self.value = self.addOutPort("value")

    def timeAdvance(self):
        if self.state["item"] is not None:
            return 0.0
        return INFINITY

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

    def intTransition(self):
        self.state["item"] = None
        self.state["num"] = self.next()
        return self.state

    def outputFnc(self):
        if self.state["item"] is not None:
            return {
                self.random: self.state["num"],
                self.value: self.state["item"]
            }
        return {}

    def next(self):
        """Obtains a random variable w.r.t. this block's config.

        Do note that this updates the gen object, meaning it must only be called
        in functions that may update the state.
        """
        return self.dist(U01(self.gen), *self.args)
