import autograd.numpy as anp
import numpy as np

from pymoo.model.problem import Problem


class Truss2D(Problem):

    def __init__(self):
        super().__init__(n_var=3, n_obj=2, n_constr=1, type_var=anp.double)

        self.Amax = 0.01
        self.Smax = 1e5

        self.xl = anp.array([0.0, 0.0, 1.0])
        self.xu = anp.array([self.Amax, self.Amax, 3.0])

    def _evaluate(self, x, out, *args, **kwargs):

        # variable names for convenient access
        x1 = x[:, 0]
        x2 = x[:, 1]
        y = x[:, 2]

        # first objectives
        f1 = x1 * anp.sqrt(16 + anp.square(y)) + x2 * anp.sqrt((1 + anp.square(y)))

        # measure which are needed for the second objective
        sigma_ac = 20 * anp.sqrt(16 + anp.square(y)) / (y * x1)
        sigma_bc = 80 * anp.sqrt(1 + anp.square(y)) / (y * x2)

        # take the max
        f2 = anp.max(anp.column_stack((sigma_ac, sigma_bc)), axis=1)

        # define a constraint
        g1 = f2 - self.Smax

        out["F"] = anp.column_stack([f1, f2])
        out["G"] = g1

    def _calc_pareto_front(self, *args, **kwargs):

        T = 2 * np.sqrt(5) * self.Amax

        # Part A - before transition point T

        f1 = np.linspace(400 / self.Smax, T, 1000)
        f2 = 400 / f1
        part_a = np.column_stack([f1, f2])

        # Part B - after transition point T

        def calc_y(V):
            return np.sqrt(3200 * V**2 + 40 * V * np.sqrt(6400 * V**2 - 12) - 4)

        def calc_SV(y):
            return (4 + y**2) / (0.01 * y)

        def calc_x1(y):
            return 0.0025 * np.sqrt((16 + y**2) / (1 + y**2))

        def calc_S(V):
            y = calc_y(V)
            SV = calc_SV(y)
            S = SV / V
            return S

        y = 3
        x1 = calc_x1(y)
        x2 = 0.01

        V_min = T
        V_max = self.evaluate(np.array([x1, x2, y]), return_values_of=["F"])[0]

        f1 = np.linspace(V_min, V_max, 100)
        f2 = calc_S(f1)
        part_b = np.column_stack([f1, f2])

        pf = np.row_stack([part_a, part_b])

        return pf


