from zoofs.baseoptimizationalgorithm import BaseOptimizationAlgorithm
import numpy as np
import time
import warnings


class GreyWolfOptimization(BaseOptimizationAlgorithm):
    def __init__(
        self,
        objective_function,
        n_iteration: int = 1000,
        timeout: int = None,
        population_size=50,
        method=1,
        minimize=True,
        logger=None,
        **kwargs,
    ):
        """
        Parameters
        ----------
        objective_function : user made function of the signature 'func(model,X_train,y_train,X_test,y_test)'
            The function must return a value, that needs to be minimized/maximized.

        n_iteration : int, default=1000
            Number of time the Optimization algorithm will run

        timeout: int = None
            Stop operation after the given number of second(s).
            If argument is set to None, the operation is executed without time limitation and n_iteration is followed

        population_size : int, default=50
            Total size of the population

        method : {1, 2}, default=1
            Choose the between the two methods of grey wolf optimization

        minimize : bool, default=True
            Defines if the objective value is to be maximized or minimized

        logger: Logger or None, optional (default=None)
            - accepts `logging.Logger` instance.

        **kwargs
            Any extra keyword argument for objective_function

        Attributes
        ----------
        best_feature_list : ndarray of shape (n_features)
            list of features with the best result of the entire run
        """
        super().__init__(
            objective_function, n_iteration, timeout, population_size, minimize, logger, **kwargs
        )
        self.method = method

    def _check_params(self, model, x_train, y_train, x_valid, y_valid, method=1):
        super()._check_params(model, x_train, y_train, x_valid, y_valid)
        if method not in [1, 2]:
            raise ValueError("method accepts only 1,2 ")

    def fit(self, model, X_train, y_train, X_valid, y_valid, verbose=True):
        """
        Parameters
        ----------
        model : machine learning model's object
           machine learning model's object

        X_train : pandas.core.frame.DataFrame of shape (n_samples, n_features)
           Training input samples to be used for machine learning model

        y_train : pandas.core.frame.DataFrame or pandas.core.series.Series of shape (n_samples)
           The target values (class labels in classification, real numbers in regression).

        X_valid : pandas.core.frame.DataFrame of shape (n_samples, n_features)
           Validation input samples

        y_valid : pandas.core.frame.DataFrame or pandas.core.series.Series of shape (n_samples)
            The target values (class labels in classification, real numbers in regression).

        verbose : bool,default=True
             Print results for iterations
        """
        self._check_params(model, X_train, y_train, X_valid, y_valid, self.method)

        self.feature_score_hash = {}
        self.feature_list = np.array(list(X_train.columns))
        self.best_results_per_iteration = {}
        self.best_score = np.inf
        self.best_dim = np.ones(X_train.shape[1])

        self.initialize_population(X_train)

        self.best_score_dimension = np.ones(X_train.shape[1])

        self.alpha_wolf_dimension, self.alpha_wolf_fitness = np.ones(X_train.shape[1]), np.inf
        self.beta_wolf_dimension, self.beta_wolf_fitness = np.ones(X_train.shape[1]), np.inf
        self.delta_wolf_dimension, self.delta_wolf_fitness = np.ones(X_train.shape[1]), np.inf

        if self.timeout is not None:
            timeout_upper_limit = time.time() + self.timeout
        else:
            timeout_upper_limit = time.time()
        for i in range(self.n_iteration):
            if (self.timeout is not None) & (time.time() > timeout_upper_limit):
                warnings.warn("Timeout occured")
                break
            a = 2 - 2 * ((i + 1) / self.n_iteration)

            self.fitness_scores = self._evaluate_fitness(model, X_train, y_train, X_valid, y_valid)

            self.iteration_objective_score_monitor(i)

            top_three_fitness_indexes = np.argsort(self.fitness_scores)[:3]

            for fit, dim in zip(
                np.array(self.fitness_scores)[top_three_fitness_indexes],
                self.individuals[top_three_fitness_indexes],
            ):
                if fit < self.alpha_wolf_fitness:
                    self.delta_wolf_fitness = self.beta_wolf_fitness
                    self.beta_wolf_fitness = self.alpha_wolf_fitness
                    self.alpha_wolf_fitness = fit

                    self.delta_wolf_dimension = self.beta_wolf_dimension
                    self.beta_wolf_dimension = self.alpha_wolf_dimension
                    self.alpha_wolf_dimension = dim
                    continue

                if (fit > self.alpha_wolf_fitness) & (fit < self.beta_wolf_fitness):
                    self.delta_wolf_fitness = self.beta_wolf_fitness
                    self.beta_wolf_fitness = fit

                    self.delta_wolf_dimension = self.beta_wolf_dimension
                    self.beta_wolf_dimension = dim
                    continue

                if (fit > self.beta_wolf_fitness) & (fit < self.delta_wolf_fitness):
                    self.delta_wolf_fitness = fit
                    self.delta_wolf_dimension = dim

            if (self.method == 1) | (self.method == 2):
                C1 = 2 * np.random.random((self.population_size, X_train.shape[1]))
                A1 = 2 * a * np.random.random((self.population_size, X_train.shape[1])) - a
                d_alpha = abs(C1 * self.alpha_wolf_dimension - self.individuals)

                C2 = 2 * np.random.random((self.population_size, X_train.shape[1]))
                A2 = 2 * a * np.random.random((self.population_size, X_train.shape[1])) - a
                d_beta = abs(C2 * self.beta_wolf_dimension - self.individuals)

                C3 = 2 * np.random.random((self.population_size, X_train.shape[1]))
                A3 = 2 * a * np.random.random((self.population_size, X_train.shape[1])) - a
                d_delta = abs(C3 * self.delta_wolf_dimension - self.individuals)

            if self.method == 2:
                X1 = abs(self.alpha_wolf_dimension - A1 * d_alpha)
                X2 = abs(self.beta_wolf_dimension - A2 * d_beta)
                X3 = abs(self.delta_wolf_dimension - A3 * d_delta)
                self.individuals = np.where(
                    np.random.uniform(size=(self.population_size, X_train.shape[1]))
                    <= self.sigmoid((X1 + X2 + X3) / 3),
                    1,
                    0,
                )

            if self.method == 1:
                Y1 = np.where(
                    (
                        self.alpha_wolf_dimension
                        + np.where(
                            self.sigmoid(A1 * d_alpha)
                            > np.random.uniform(size=(self.population_size, X_train.shape[1])),
                            1,
                            0,
                        )
                    )
                    >= 1,
                    1,
                    0,
                )
                Y2 = np.where(
                    (
                        self.beta_wolf_dimension
                        + np.where(
                            self.sigmoid(A1 * d_beta)
                            > np.random.uniform(size=(self.population_size, X_train.shape[1])),
                            1,
                            0,
                        )
                    )
                    >= 1,
                    1,
                    0,
                )
                Y3 = np.where(
                    (
                        self.delta_wolf_dimension
                        + np.where(
                            self.sigmoid(A1 * d_delta)
                            > np.random.uniform(size=(self.population_size, X_train.shape[1])),
                            1,
                            0,
                        )
                    )
                    >= 1,
                    1,
                    0,
                )
                r = np.random.uniform(size=(self.population_size, X_train.shape[1]))
                self.individuals[r < (1 / 3)] = Y1[r < (1 / 3)]
                self.individuals[(r >= (1 / 3)) & (r < (2 / 3))] = Y2[
                    (r >= (1 / 3)) & (r < (2 / 3))
                ]
                self.individuals[r >= (2 / 3)] = Y3[r >= (2 / 3)]

            self.verbose_results(verbose, i)
            self.best_feature_list = list(self.feature_list[np.where(self.best_dim)[0]])
        return self.best_feature_list
