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

[((29, 10, 29, 52), 'pymortests.fixtures.parameter.parameters_of_type', 'parameters_of_type', ({(29, 29, 29, 45): 'f.parameter_type', (29, 47, 29, 51): '4711'}, {}), '(f.parameter_type, 4711)', False, 'from pymortests.fixtures.parameter import parameters_of_type\n'), ((62, 4, 62, 30), 'pymortests.pickling.assert_picklable', 'assert_picklable', ({(62, 21, 62, 29): 'function'}, {}), '(function)', False, 'from pymortests.pickling import assert_picklable, assert_picklable_without_dumps_function\n'), ((66, 4, 66, 63), 'pymortests.pickling.assert_picklable_without_dumps_function', 'assert_picklable_without_dumps_function', ({(66, 44, 66, 62): 'picklable_function'}, {}), '(picklable_function)', False, 'from pymortests.pickling import assert_picklable, assert_picklable_without_dumps_function\n'), ((72, 10, 72, 50), 'pymortests.fixtures.parameter.parameters_of_type', 'parameters_of_type', ({(72, 29, 72, 45): 'f.parameter_type', (72, 47, 72, 49): '47'}, {}), '(f.parameter_type, 47)', False, 'from pymortests.fixtures.parameter import parameters_of_type\n'), ((73, 15, 73, 43), 'pymortests.fixtures.function.function_argument', 'function_argument', ({(73, 33, 73, 34): 'f', (73, 36, 73, 38): '(10)', (73, 40, 73, 42): '(42)'}, {}), '(f, 10, 42)', False, 'from pymortests.fixtures.function import function, picklable_function, function_argument\n'), ((31, 14, 31, 46), 'pymortests.fixtures.function.function_argument', 'function_argument', ({(31, 32, 31, 33): 'f', (31, 35, 31, 40): 'count', (31, 42, 31, 45): '454'}, {}), '(f, count, 454)', False, 'from pymortests.fixtures.function import function, picklable_function, function_argument\n'), ((38, 12, 38, 40), 'numpy.linspace', 'np.linspace', (), '', True, 'import numpy as np\n'), ((39, 15, 39, 54), 'pymor.functions.basic.ConstantFunction', 'ConstantFunction', (), '', False, 'from pymor.functions.basic import ConstantFunction, GenericFunction\n'), ((57, 9, 57, 38), 'pytest.raises', 'pytest.raises', ({(57, 23, 57, 37): 'AssertionError'}, {}), '(AssertionError)', False, 'import pytest\n'), ((58, 8, 58, 38), 'pymor.functions.basic.ConstantFunction', 'ConstantFunction', (), '', False, 'from pymor.functions.basic import ConstantFunction, GenericFunction\n'), ((71, 15, 71, 23), 'pymor.core.pickle.dumps', 'dumps', ({(71, 21, 71, 22): 'f'}, {}), '(f)', False, 'from pymor.core.pickle import dumps, loads\n'), ((40, 21, 40, 60), 'pymor.functions.basic.ConstantFunction', 'ConstantFunction', (), '', False, 'from pymor.functions.basic import ConstantFunction, GenericFunction\n'), ((53, 13, 53, 42), 'pytest.raises', 'pytest.raises', ({(53, 27, 53, 41): 'AssertionError'}, {}), '(AssertionError)', False, 'import pytest\n'), ((55, 13, 55, 42), 'pytest.raises', 'pytest.raises', ({(55, 27, 55, 41): 'AssertionError'}, {}), '(AssertionError)', False, 'import pytest\n'), ((42, 24, 42, 63), 'pymor.functions.basic.ConstantFunction', 'ConstantFunction', (), '', False, 'from pymor.functions.basic import ConstantFunction, GenericFunction\n'), ((54, 19, 54, 57), 'pymor.functions.basic.ConstantFunction', 'ConstantFunction', (), '', False, 'from pymor.functions.basic import ConstantFunction, GenericFunction\n'), ((56, 19, 56, 53), 'pymor.functions.basic.ConstantFunction', 'ConstantFunction', (), '', False, 'from pymor.functions.basic import ConstantFunction, GenericFunction\n'), ((41, 47, 41, 69), 'numpy.zeros', 'np.zeros', ({(41, 56, 41, 68): 'X.shape[:-1]'}, {}), '(X.shape[:-1])', True, 'import numpy as np\n'), ((45, 37, 45, 49), 'numpy.zeros', 'np.zeros', ({(45, 46, 45, 48): '()'}, {}), '(())', True, 'import numpy as np\n'), ((43, 50, 43, 71), 'numpy.ones', 'np.ones', ({(43, 58, 43, 70): 'X.shape[:-1]'}, {}), '(X.shape[:-1])', True, 'import numpy as np\n')]