# Copyright (c) 2020, Anthony Nouy, Erwan Grelier
# This file is part of tensap (tensor approximation package).

# tensap is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# tensap is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.

# You should have received a copy of the GNU Lesser General Public License
# along with tensap.  If not, see <https://www.gnu.org/licenses/>.

"""
Module multivariate_functions_benchmark

"""

import sys
import numpy as np

sys.path.insert(0, "./../../../")
import tensap  # noqa: E402


def multivariate_functions_benchmark(case, *args):
    """
    Generate different functions used to benchmark the package.

    Parameters
    ----------
    case : str
        The name of the function. Can be 'borehole', 'ishigami',
        'sin_of_a_sum', 'linear_additive', 'linear_rank_one',
        'quadratic_rank_one', 'canonical_rank_2', 'mixture', 'field',
        'oscillatory', 'product_peak', 'corner_peak', 'gaussian', 'continuous',
        'discontinuous', 'henon_heiles', 'sobol', 'anisotropic', 'polynomial'.
    *args : tuple
        Parameters of the function.

    Raises
    ------
    NotImplementedError
        If the function is not implemented.

    Returns
    -------
    function
        The asked function.
    tensap.RandomVector
        Input random variables.

    """
    if case == "borehole":
        X = np.empty(8, dtype=object)
        X[0] = tensap.NormalRandomVariable(0.1, 0.0161812)
        X[1] = tensap.NormalRandomVariable(0, 1)
        X[2] = tensap.UniformRandomVariable(63070, 115600)
        X[3] = tensap.UniformRandomVariable(990, 1110)
        X[4] = tensap.UniformRandomVariable(63.1, 116)
        X[5] = tensap.UniformRandomVariable(700, 820)
        X[6] = tensap.UniformRandomVariable(1120, 1680)
        X[7] = tensap.UniformRandomVariable(9855, 12045)
        X = tensap.RandomVector(X)

        def fun(x):
            return (
                2
                * np.pi
                * x[:, 2]
                * (x[:, 3] - x[:, 5])
                / (
                    np.log(np.exp(7.71 + 1.0056 * x[:, 1]) / x[:, 0])
                    * (
                        1
                        + 2
                        * x[:, 6]
                        * x[:, 2]
                        / np.log(np.exp(7.71 + 1.0056 * x[:, 1]) / x[:, 0])
                        / x[:, 0] ** 2
                        / x[:, 7]
                        + x[:, 2] / x[:, 4]
                    )
                )
            )

    elif case == "ishigami":
        if len(args) == 0:
            d = 3
        else:
            d = args[0]
        if len(args) <= 1:
            a = 7
        else:
            a = args[1]
        if len(args) <= 2:
            b = 0.1
        else:
            b = args[2]

        X = tensap.RandomVector(tensap.UniformRandomVariable(-np.pi, np.pi), d)

        def fun(x):
            return (
                np.sin(x[:, 0])
                + a * np.sin(x[:, 1]) ** 2
                + b * x[:, 2] ** 4 * np.sin(x[:, 0])
            )

    elif case == "sin_of_asum":
        if len(args) == 0:
            d = 3
        else:
            d = args[0]
        if len(args) <= 1:

            def fun(x):
                return np.sin(np.sum(x, 1))

        else:

            def fun(x):
                return np.sin(np.matmul(x, np.reshape(args[1], (-1, 1))))

        X = tensap.RandomVector(tensap.UniformRandomVariable(), d)

    elif case == "linear_additive":
        if len(args) == 0:
            d = 3
        else:
            d = args[0]
        if len(args) <= 1:

            def fun(x):
                return np.sum(x, 1)

        else:

            def fun(x):
                return np.matmul(x, np.reshape(args[1], (-1, 1)))

        X = tensap.RandomVector(tensap.UniformRandomVariable(), d)

    elif case == "linear_rank_one":
        if len(args) == 0:
            d = 3
        else:
            d = args[0]
        if len(args) <= 1:
            w = np.zeros((1, d))
        else:
            w = np.reshape(args[1], (1, -1))

        X = tensap.RandomVector(tensap.UniformRandomVariable(), d)

        def fun(x):
            return np.prod(np.tile(w, (np.shape(x)[0], 1)) + 2 * x, 1)

    elif case == "quadratic_rank_one":
        d = args[0]
        c = args[1]

        X = tensap.RandomVector(tensap.UniformRandomVariable(), d)

        def fun(x):
            return np.prod(
                np.tile(c[0, :], (np.shape(x)[0], 1))
                + np.tile(c[1, :], (np.shape(x)[0], 1)) * x
                + np.tile(c[2, :], (np.shape(x)[0], 1)) * x**2,
                1,
            )

    elif case == "canonical_rank_2":
        if len(args) == 0:
            d = 3
        else:
            d = args[0]

        X = tensap.RandomVector(tensap.UniformRandomVariable(), d)

        def fun(x):
            return x[:, 0] * x[:, 1] * x[:, 2] + x[:, 0] ** 2 + x[:, 1]

    elif case == "mixture":
        if len(args) == 0:
            d = 6
        else:
            d = args[0]

        X = tensap.RandomVector(tensap.UniformRandomVariable(), d)

        def fun(x):
            return (
                np.sin(x[:, 0] + x[:, 3]) * np.exp(x[:, 4]) * x[:, 5]
                + np.sin(x[:, 2] * x[:, 3]) * x[:, 5]
            )

    elif case == "field":
        if len(args) == 0:
            d = 6
        else:
            d = args[0]

        X = tensap.RandomVector(tensap.UniformRandomVariable(0, 1), d)

        def fun(x):
            return (
                1
                + np.cos(x[:, 0]) * x[:, 1]
                + np.sin(x[:, 0]) * x[:, 2]
                + np.exp(x[:, 0]) * x[:, 3]
                + 1 / (x[:, 0] + 1) * x[:, 4]
                + 1 / (2 * x[:, 0] + 3) * x[:, 5]
            )

    elif case in (
        "oscillatory",
        "product_peak",
        "corner_peak",
        "gaussian",
        "continuous",
        "discontinuous",
    ):
        if len(args) == 0:
            d = 10
        else:
            d = args[0]

        w = np.random.rand(1, d)
        c = np.random.rand(1, d)

        X = tensap.RandomVector(tensap.UniformRandomVariable(0, 1), d)

        if case == "oscillatory":

            def fun(x):
                return np.cos(w[0] * 2 * np.pi + np.matmul(x, np.transpose(c)))

        elif case == "product_peak":

            def fun(x):
                return 1 / np.prod(
                    np.tile(x ** (-2), (np.shape(x)[0], 1))
                    + (x + np.tile(w, (np.shape(x)[0], 1))) ** 2,
                    1,
                )

        elif case == "corner_peak":
            b = 185
            e = 2
            c = c * b / d**e / np.sum(c)

            def fun(x):
                return (1 + np.matmul(x, np.transpose(c))) ** (-d - 1)

        elif case == "gaussian":

            def fun(x):
                return np.matmul(
                    np.exp(
                        -((x - np.tile(w, (np.shape(x)[0], 1))) ** 2),
                        np.transpose(c) ** 2,
                    )
                )

        elif case == "continuous":

            def fun(x):
                return np.matmul(
                    np.exp(
                        -np.abs(x - np.tile(w, (np.shape(x)[0], 1))),
                        np.transpose(c) ** 2,
                    )
                )

        elif case == "discontinuous":
            raise NotImplementedError("Function not implemented.")

    elif case == "henon_heiles":
        if len(args) == 0:
            d = 3
        else:
            d = args[0]

        X = tensap.RandomVector(tensap.NormalRandomVariable(), d)

        def fun(x):
            return (
                0.5 * np.sum(x**2, 1)
                + 0.2 * np.sum(x[:, :-1] * x[:, 1:] ** 2 - x[:, :-1] ** 3, 1)
                + 0.2**2 / 16 * np.sum((x[:, :-1] ** 2 + x[:, 1:] ** 2) ** 2, 1)
            )

    elif case == "sobol":
        if len(args) == 0:
            d = 8
            a = np.array([1, 2, 5, 10, 20, 50, 100, 500])
        else:
            d = args[0]
            if len(args) <= 1:
                a = 2 ** np.arange(0, d)
            else:
                a = np.array(args[1])

        X = tensap.RandomVector(tensap.UniformRandomVariable(0, 1), d)

        def fun(x):
            return np.prod(
                (np.abs(4 * x - 2) + np.tile(a, (np.shape(x)[0], 1)))
                / np.tile(1 + a, (np.shape(x)[0], 1)),
                1,
            )

    elif case == "anisotropic":
        if len(args) == 0:
            d = 16
        else:
            d = args[0]

        X = tensap.RandomVector(tensap.UniformRandomVariable(0, 1), d)

        def fun(x):
            return x[:, 2] * np.sin(x[:, 3] + x[:, 15])

    elif case == "polynomial":
        if len(args) == 0:
            d = 16
        else:
            d = args[0]
        if len(args) <= 1:
            q = 2
        else:
            q = args[1]

        X = tensap.RandomVector(tensap.UniformRandomVariable(0, 1), d)

        def fun(x):
            return 1 / (2 ** np.shape(x)[1]) * np.prod(3 * x**q + 1, 1)

    else:
        raise NotImplementedError("Bad function name.")

    return fun, X
