# Source code for pypdevsbbl.generic.math

    #  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:: ../paper/images/blocks/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:: ../paper/images/blocks/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:: ../paper/images/blocks/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:: ../paper/images/blocks/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:: ../paper/images/blocks/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:: ../paper/images/blocks/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)