import numpy as np
import matplotlib.pyplot as plt


class History:
    """
    This class is used to track the history of the optimizer. It tracks the global best, current best, and the time of
    the epoch.
    Also, it provides methods to visualize the history of the optimizer.

    """

    def __init__(self, **kwargs):
        self.optimizer = kwargs['optimizer']
        self.epoch = kwargs['epoch']
        self.population_size = kwargs['population_size']

        self.global_best_list = np.zeros((self.epoch, self.optimizer.dimensions))  # coordinates of the global best
        self.global_best_fitness_list = np.zeros(self.epoch)  # fitness of the global best
        self.current_best_list = np.zeros((self.epoch, self.optimizer.dimensions))  # coordinates of the current best
        self.current_best_fitness_list = np.zeros(self.epoch)  # fitness of the current best
        self.epoch_time_list = np.zeros(self.epoch)  # time of the current epoch

    def update_history(self, epoch, epoch_time):
        """
        This method updates the history of the optimizer
        """
        self.global_best_list[epoch] = self.optimizer.get_best_solution()
        self.global_best_fitness_list[epoch] = self.optimizer.get_best_score()
        self.current_best_list[epoch] = self.optimizer.get_best_solution()
        self.current_best_fitness_list[epoch] = self.optimizer.get_best_score()
        self.epoch_time_list[epoch] = epoch_time

    def get_history(self) -> dict:
        """
        This method returns the history of the optimizer
        """
        return {
            'global_best_list': self.global_best_list,
            'global_best_fitness_list': self.global_best_fitness_list,
            'current_best_list': self.current_best_list,
            'current_best_fitness_list': self.current_best_fitness_list,
            'epoch_time_list': self.epoch_time_list
        }

    def get_global_best(self) -> tuple:
        """
        This method returns the global best
        """
        return self.global_best_list[-1], self.global_best_fitness_list[-1]

    def get_current_best(self) -> tuple:
        """
        This method returns the current best
        """
        return self.current_best_list[-1], self.current_best_fitness_list[-1]

    @staticmethod
    def visualize(func):
        """
        Decorator to visualize the history of the optimizer

        :param func: function to be decorated
        """

        def wrapper(*args, **kwargs):
            plt.figure(figsize=(12, 8))
            func(*args, **kwargs)
            plt.legend()
            plt.show()

        return wrapper

    @visualize
    def visualize_fitness(self):
        """
        This method visualizes the history of the optimizer
        """
        plt.plot(self.global_best_fitness_list, label='Global Best')
        plt.plot(self.current_best_fitness_list, label='Current Best')
        plt.xlabel('Epochs')
        plt.ylabel('Fitness')
        plt.title('Fitness vs Epochs')

    @visualize
    def visualize_time(self):
        """
        This method visualizes the history of the optimizer
        """
        plt.plot(self.epoch_time_list, label='Epoch Time')
        plt.xlabel('Epochs')
        plt.ylabel('Time')
        plt.title('Time vs Epochs')

    def is_early_stopping(self, current_epoch, early_stopping: int):
        """
        This method checks if the early stopping condition is met

        :param current_epoch: current epoch
        :param early_stopping: number of epochs to wait before stopping the optimization

        :return: boolean
        """
        if early_stopping is not None and current_epoch >= early_stopping:
            last_scores = self.global_best_fitness_list[current_epoch - early_stopping + 1:current_epoch + 1]
            if len(last_scores) == early_stopping and np.all(last_scores == last_scores[0]):
                return True
        return False
