import matplotlib.pyplot as plt
import numpy as np

from pymoo.visualization.fitness_landscape import FitnessLandscape
from pymoo.visualization.video.callback_video import AnimationCallback


class TwoVariablesOneObjectiveVisualization(AnimationCallback):

    def __init__(self,
                 n_samples_for_surface=10000,
                 **kwargs):
        super().__init__(**kwargs)
        self.last_pop = None
        self.n_samples_for_surface = n_samples_for_surface

    def do(self, problem, algorithm):

        # check whether the visualization can be done or not - throw exception or simply do nothing
        if problem.n_var != 2 or problem.n_obj != 1:
            raise Exception("This visualization can only be used for problems with two variables and one objective!")

        # draw the problem surface
        FitnessLandscape(problem, _type="contour", kwargs_contour=dict(alpha=0.5)).do()

        # get the population
        pop = algorithm.pop

        X, F, CV = pop.get("X", "F", "CV")
        plt.scatter(X[:, 0], X[:, 1], color="blue", marker="o", s=70)

        is_new = np.full(len(pop), True)
        if self.last_pop is not None:
            for k, ind in enumerate(pop):
                if ind in self.last_pop:
                    is_new[k] = False

        # plot the new population
        if is_new.sum() > 0:
            X, F, CV = pop[is_new].get("X", "F", "CV")
            plt.scatter(X[:, 0], X[:, 1], color="red", marker="*", s=70)

        if hasattr(algorithm, "off") and algorithm.off is not None:
            X, F, CV = algorithm.off.get("X", "F", "CV")
            plt.scatter(X[:, 0], X[:, 1], color="purple", marker="*", s=40)

        xl, xu = problem.bounds()
        plt.xlim(xl[0], xu[0])
        plt.ylim(xl[1], xu[1])

        plt.title(f"Generation: {algorithm.n_gen}")
        plt.legend()

        # store the current population as the last
        self.last_pop = set(pop)
