# This file is part of the pyMOR project (http://www.pymor.org).
# Copyright 2013-2017 pyMOR developers and contributors. All rights reserved.
# License: BSD 2-Clause License (http://opensource.org/licenses/BSD-2-Clause)

import numpy as np
import pytest

from pymor.core.pickle import dumps, loads
from pymor.functions.basic import ConstantFunction, GenericFunction
from pymortests.fixtures.function import function, picklable_function, function_argument
from pymortests.fixtures.parameter import parameters_of_type
from pymortests.pickling import assert_picklable, assert_picklable_without_dumps_function


# monkey np.testing.assert_allclose to behave the same as np.allclose
# for some reason, the default atol of np.testing.assert_allclose is 0
# while it is 1e-8 for np.allclose

real_assert_allclose = np.testing.assert_allclose


def monkey_allclose(a, b, rtol=1.e-5, atol=1.e-8):
    real_assert_allclose(a, b, rtol=rtol, atol=atol)
np.testing.assert_allclose = monkey_allclose


def test_evaluate(function):
    f = function
    mus = parameters_of_type(f.parameter_type, 4711)
    for count in [0, 1, 5, (0, 1), (2, 2, 2)]:
        arg = function_argument(f, count, 454)
        result = f.evaluate(arg, next(mus))
        assert result.shape == arg.shape[:-1] + f.shape_range


def test_lincomb_function():
    for steps in (1, 10):
        x = np.linspace(0, 1, num=steps)
        zero = ConstantFunction(0.0, dim_domain=steps)
        for zero in (ConstantFunction(0.0, dim_domain=steps),
                     GenericFunction(lambda X: np.zeros(X.shape[:-1]), dim_domain=steps)):
            for one in (ConstantFunction(1.0, dim_domain=steps),
                        GenericFunction(lambda X: np.ones(X.shape[:-1]), dim_domain=steps), 1.0):
                add = (zero + one) + 0
                sub = (zero - one) + np.zeros(())
                neg = - zero
                assert np.allclose(sub(x), [-1])
                assert np.allclose(add(x), [1.0])
                assert np.allclose(neg(x), [0.0])
                (repr(add), str(add), repr(one), str(one))  # just to cover the respective special funcs too
                mul = neg * 1.
                assert np.allclose(mul(x), [0.0])
        with pytest.raises(AssertionError):
            zero + ConstantFunction(dim_domain=steps + 1)
        with pytest.raises(AssertionError):
            zero * ConstantFunction(dim_domain=steps)
    with pytest.raises(AssertionError):
        ConstantFunction(dim_domain=0)


def test_pickle(function):
    assert_picklable(function)


def test_pickle_without_dumps_function(picklable_function):
    assert_picklable_without_dumps_function(picklable_function)


def test_pickle_by_evaluation(function):
    f = function
    f2 = loads(dumps(f))
    mus = parameters_of_type(f.parameter_type, 47)
    for arg in function_argument(f, 10, 42):
        mu = next(mus)
        assert np.all(f.evaluate(arg, mu) == f2.evaluate(arg, mu))

[((1045, 1087), 'pymortests.fixtures.parameter.parameters_of_type', 'parameters_of_type', (['f.parameter_type', '(4711)'], {}), '(f.parameter_type, 4711)\n', (1063, 1087), False, 'from pymortests.fixtures.parameter import parameters_of_type\n'), ((2526, 2552), 'pymortests.pickling.assert_picklable', 'assert_picklable', (['function'], {}), '(function)\n', (2542, 2552), False, 'from pymortests.pickling import assert_picklable, assert_picklable_without_dumps_function\n'), ((2619, 2678), 'pymortests.pickling.assert_picklable_without_dumps_function', 'assert_picklable_without_dumps_function', (['picklable_function'], {}), '(picklable_function)\n', (2658, 2678), False, 'from pymortests.pickling import assert_picklable, assert_picklable_without_dumps_function\n'), ((2774, 2814), 'pymortests.fixtures.parameter.parameters_of_type', 'parameters_of_type', (['f.parameter_type', '(47)'], {}), '(f.parameter_type, 47)\n', (2792, 2814), False, 'from pymortests.fixtures.parameter import parameters_of_type\n'), ((2830, 2858), 'pymortests.fixtures.function.function_argument', 'function_argument', (['f', '(10)', '(42)'], {}), '(f, 10, 42)\n', (2847, 2858), False, 'from pymortests.fixtures.function import function, picklable_function, function_argument\n'), ((1149, 1181), 'pymortests.fixtures.function.function_argument', 'function_argument', (['f', 'count', '(454)'], {}), '(f, count, 454)\n', (1166, 1181), False, 'from pymortests.fixtures.function import function, picklable_function, function_argument\n'), ((1357, 1385), 'numpy.linspace', 'np.linspace', (['(0)', '(1)'], {'num': 'steps'}), '(0, 1, num=steps)\n', (1368, 1385), True, 'import numpy as np\n'), ((1401, 1440), 'pymor.functions.basic.ConstantFunction', 'ConstantFunction', (['(0.0)'], {'dim_domain': 'steps'}), '(0.0, dim_domain=steps)\n', (1417, 1440), False, 'from pymor.functions.basic import ConstantFunction, GenericFunction\n'), ((2423, 2452), 'pytest.raises', 'pytest.raises', (['AssertionError'], {}), '(AssertionError)\n', (2436, 2452), False, 'import pytest\n'), ((2462, 2492), 'pymor.functions.basic.ConstantFunction', 'ConstantFunction', ([], {'dim_domain': '(0)'}), '(dim_domain=0)\n', (2478, 2492), False, 'from pymor.functions.basic import ConstantFunction, GenericFunction\n'), ((2754, 2762), 'pymor.core.pickle.dumps', 'dumps', (['f'], {}), '(f)\n', (2759, 2762), False, 'from pymor.core.pickle import dumps, loads\n'), ((1462, 1501), 'pymor.functions.basic.ConstantFunction', 'ConstantFunction', (['(0.0)'], {'dim_domain': 'steps'}), '(0.0, dim_domain=steps)\n', (1478, 1501), False, 'from pymor.functions.basic import ConstantFunction, GenericFunction\n'), ((2227, 2256), 'pytest.raises', 'pytest.raises', (['AssertionError'], {}), '(AssertionError)\n', (2240, 2256), False, 'import pytest\n'), ((2329, 2358), 'pytest.raises', 'pytest.raises', (['AssertionError'], {}), '(AssertionError)\n', (2342, 2358), False, 'import pytest\n'), ((1618, 1657), 'pymor.functions.basic.ConstantFunction', 'ConstantFunction', (['(1.0)'], {'dim_domain': 'steps'}), '(1.0, dim_domain=steps)\n', (1634, 1657), False, 'from pymor.functions.basic import ConstantFunction, GenericFunction\n'), ((2277, 2315), 'pymor.functions.basic.ConstantFunction', 'ConstantFunction', ([], {'dim_domain': '(steps + 1)'}), '(dim_domain=steps + 1)\n', (2293, 2315), False, 'from pymor.functions.basic import ConstantFunction, GenericFunction\n'), ((2379, 2413), 'pymor.functions.basic.ConstantFunction', 'ConstantFunction', ([], {'dim_domain': 'steps'}), '(dim_domain=steps)\n', (2395, 2413), False, 'from pymor.functions.basic import ConstantFunction, GenericFunction\n'), ((1550, 1572), 'numpy.zeros', 'np.zeros', (['X.shape[:-1]'], {}), '(X.shape[:-1])\n', (1558, 1572), True, 'import numpy as np\n'), ((1833, 1845), 'numpy.zeros', 'np.zeros', (['()'], {}), '(())\n', (1841, 1845), True, 'import numpy as np\n'), ((1709, 1730), 'numpy.ones', 'np.ones', (['X.shape[:-1]'], {}), '(X.shape[:-1])\n', (1716, 1730), True, 'import numpy as np\n')]