import numpy as np


class VLMOP2():
    '''
    Van Veldhuizen, David A., and Gary B. Lamont. "Multiobjective evolutionary algorithm test suites." Proceedings of the 1999 ACM symposium on Applied computing. 1999.
    '''
    def __init__(self, n_var=2):
        self.n_var = n_var
        self.n_obj = 2
        self.xl = 0.0
        self.xu = 1.0
        self.vlmop2_xl = -2.0 * np.ones(n_var)
        self.vlmop2_xu = 2.0 * np.ones(n_var)

    def _evaluate(self, x):
        n = self.n_var

        f1 = 1 - np.exp(-np.sum((x - 1 / np.sqrt(n)) ** 2, axis=1))
        f2 = 1 - np.exp(-np.sum((x + 1 / np.sqrt(n)) ** 2, axis=1))

        return np.column_stack([f1, f2])
    
    def _calc_pareto_front(self, n_pareto_points=100):
        n = self.n_var

        x = np.linspace(-1 / np.sqrt(n), 1 / np.sqrt(n), n_pareto_points)
        x_all = np.column_stack([x] * n)

        return self._evaluate(x_all)
    
    def __call__(self, x, *args, **kwargs):
        assert x.ndim == 2
        assert x.shape[1] == self.n_var
        assert np.all(x >= self.xl) and np.all(x <= self.xu)
        # scale from [0, 1] to [vlmop2_xl, vlmop2_xu]
        x_scaled = self.vlmop2_xl + (self.vlmop2_xu - self.vlmop2_xl) * x
        fx = self._evaluate(x_scaled, *args, **kwargs)
        return fx


class VLMOP3():
    '''
    Van Veldhuizen, David A., and Gary B. Lamont. "Multiobjective evolutionary algorithm test suites." Proceedings of the 1999 ACM symposium on Applied computing. 1999.
    '''
    def __init__(self):
        self.n_var = 2
        self.n_obj = 3
        self.xl = 0.0
        self.xu = 1.0
        self.vlmop3_xl = np.array([-3.0, -3.0])
        self.vlmop3_xu = np.array([3.0, 3.0])

    def _evaluate(self, x):
        x1, x2 = x[:, 0], x[:, 1]

        f1 = 0.5 * (x1 ** 2 + x2 ** 2) + np.sin(x1 ** 2 + x2 ** 2)
        f2 = (3 * x1 - 2 * x2 + 4) ** 2 / 8 + (x1 - x2 + 1) ** 2 / 27 + 15
        f3 = 1 / (x1 ** 2 + x2 ** 2 + 1) - 1.1 * np.exp(-x1 ** 2 - x2 ** 2)

        return np.column_stack([f1, f2, f3])

    def __call__(self, x, *args, **kwargs):
        assert x.ndim == 2
        assert x.shape[1] == self.n_var
        assert np.all(x >= self.xl) and np.all(x <= self.xu)
        # scale from [0, 1] to [vlmop3_xl, vlmop3_xu]
        x_scaled = self.vlmop3_xl + (self.vlmop3_xu - self.vlmop3_xl) * x
        fx = self._evaluate(x_scaled, *args, **kwargs)
        return fx