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

"""Tests the mathematical building blocks to prove validity w.r.t. these tests.

See Also:
    :mod:`pypdevsbbl.generic.math`
"""

import unittest

from pypdevsbbl.extra import equations
from pypdevsbbl.generic.collectors import TableCollector
from pypdevsbbl.generic.generators import ConstantGenerator, SingleFire, FunctionGenerator, RandomNumberGenerator
from pypdevsbbl.generic.math import Equation, Adder, Multiplier, Differentiator, Integrator, Random
from test.generic import Model, simulate
import numpy as np

class TestMath(unittest.TestCase):
    def test_equation(self):
        """Tests the :class:`pypdevsbbl.generic.math.Adder` block."""
        nodes = [ConstantGenerator("Gen3", 3), ConstantGenerator("Gen5", 5, 3),
                 Equation("Sub", equations.SUBTRACT), TableCollector("Coll")]
        conn = [(nodes[0].output, nodes[2].y), (nodes[1].output, nodes[2].x), (nodes[2].result, nodes[3].input)]
        model = Model(nodes, conn)
        simulate(model, 8)

        assert [(0, 2), (1, 2), (2, 2), (3, 2), (4, 2), (5, 2), (6, 2), (7, 2), (8, 2)] == nodes[3].table()

    def test_adder(self):
        """Tests the :class:`pypdevsbbl.generic.math.Adder` block."""
        nodes = [SingleFire("SF-1", 1, 1), SingleFire("SF-2", 2, 2),
                 SingleFire("SF-3", 3, 3), SingleFire("SF-5", 5, 5),
                 Adder("Adder", 3), TableCollector("Collector")]
        conn = [(nodes[0].output, nodes[4].inputs[0]), (nodes[1].output, nodes[4].inputs[1]),
                (nodes[2].output, nodes[4].inputs[2]), (nodes[3].output, nodes[4].inputs[1]),
                (nodes[4].result, nodes[5].input)]
        model = Model(nodes, conn)
        simulate(model)

        assert [(3, 6), (5, 9)] == nodes[5].table()

    def test_multiplier(self):
        """Tests the :class:`pypdevsbbl.generic.math.Multiplier` block."""
        nodes = [SingleFire("SF-1", 1, 1), SingleFire("SF-2", 2, 2),
                 SingleFire("SF-3", 3, 3), SingleFire("SF-5", 5, 5),
                 Multiplier("Multiplier", 3), TableCollector("Collector")]
        conn = [(nodes[0].output, nodes[4].inputs[0]), (nodes[1].output, nodes[4].inputs[1]),
                (nodes[2].output, nodes[4].inputs[2]), (nodes[3].output, nodes[4].inputs[1]),
                (nodes[4].result, nodes[5].input)]
        model = Model(nodes, conn)
        simulate(model)

        assert [(3, 6), (5, 15)] == nodes[5].table()

    def test_differentiator(self):
        """Tests the :class:`pypdevsbbl.generic.math.Differentiator` block."""
        def f(x):
            return x ** 3 + x ** 2

        dx = 1e-3
        nodes = [FunctionGenerator("Generator", lambda _, t: f(t), dt=dx), Differentiator("dt"), TableCollector("Coll")]
        conn = [(nodes[0].output, nodes[1].input), (nodes[1].result, nodes[2].input)]
        model = Model(nodes, conn)
        simulate(model, 10)

        def der(func, pos, step, ic):
            if pos == 0:
                return ic
            return (func(pos) - func(pos - step)) / step

        exp = [(x, der(f, x, dx, 0)) for x in np.arange(0, 10 + dx, dx)]
        act = nodes[2].table()
        assert len(exp) == len(act)
        diff = [exp[x][1] - act[x][1] for x in range(round(10 / dx))]
        assert len([True for x in diff if abs(x) > dx]) == 0

    def test_integrator(self):
        """Tests the :class:`pypdevsbbl.generic.math.Integrator` block."""
        def f(x):
            return x ** 3 + x ** 2

        dx = 1e-3
        nodes = [FunctionGenerator("Generator", lambda _, t: f(t), dt=dx), Integrator("I"), TableCollector("Coll")]
        conn = [(nodes[0].output, nodes[1].input), (nodes[1].result, nodes[2].input)]
        model = Model(nodes, conn)
        simulate(model, 10)

        assert abs(np.trapz([f(x) for x in np.arange(0, 10 + dx, dx)], dx=dx) - nodes[2].table()[-1][1]) < dx

    def test_random(self):
        nodes = [ConstantGenerator("Generator"), Random("Random", seed=42), TableCollector("Coll-1"),
                 RandomNumberGenerator("RNG", seed=42), TableCollector("Coll-2")]
        conn = [(nodes[0].output, nodes[1].input), (nodes[1].random, nodes[2].input),
                (nodes[3].output, nodes[4].input)]
        model = Model(nodes, conn)
        simulate(model, 10)

        assert nodes[2].table() == nodes[4].table()

if __name__ == '__main__':
    unittest.main()
