import numpy as np
import benchmark_functions
from copy import deepcopy

"""
    GreyWolf is characterized by several parameters, limits in which he is located, objetcive function for his fitness
    dimension for number of dimensions of parameters and fitness. Specific for him are three parameters is_alpha
    is_beta and is_delta. These serve as flag, for determination wheter the wolf is one of three best solutions in
    iteration, through which every other wolf ( called omega ) updates its position
    For this status -> wheter he is alpha beta or delta, respective functions are implemented, for flag set and reset
    Omega_wolf_move method is called only on GreyWolf object wolf, which is neither alpha nor beta nor delta -> therefor
    he is omega wolf, who is trying to update its position with these three best solutions
    There are three vectors first two A and C are random vectors, each for one wolf ( alpha beta.. ) through them, X is
    computed which is a step of omega wolf in direction to each best wolf ( omega to alpha, omega to beta, omega to 
    delta). Finally, new position of omega wolf is determined by arithmetic mean of these three steps ( components of X
    vector ), then limits are checked and after that fitness is updated
"""


class GreyWolf:
    def __init__(self, limits, obj_function, dim):
        self.lim = limits
        self.obj_function = obj_function
        self.dim = dim
        self.position = np.random.uniform(self.lim[0], self.lim[1], self.dim)
        self.fitness = self.obj_function(self.position, self.dim)
        self.is_alpha = False
        self.is_beta = False
        self.is_delta = False

    def set_alpha_status(self):
        self.is_alpha = True

    def set_beta_status(self):
        self.is_beta = True

    def set_delta_status(self):
        self.is_delta = True

    def reset_alpha_status(self):
        self.is_alpha = False

    def reset_beta_status(self):
        self.is_beta = False

    def reset_delta_status(self):
        self.is_delta = False

    def omega_wolf_move(self, alpha_wolf, beta_wolf, delta_wolf, a):
        A = np.zeros(3)
        C = np.zeros(3)
        D = np.zeros(3)
        for part in range(3):
            r1 = np.random.random()
            r2 = np.random.random()
            A[part] = 2 * a * r1 - a
            C[part] = 2 * r2

        abd_wolves = [alpha_wolf, beta_wolf, delta_wolf]
        X = np.zeros([3, self.dim])

        for part in range(3):
            X[part] = self.step(abd_wolves[part], A[part], C[part])

        self.update_position(X)

    def step(self, abd_wolf, A, C):
        X = np.zeros(self.dim)
        for d in range(self.dim):
            D = np.abs(C * abd_wolf.position[d] - self.position[d])
            X[d] = abd_wolf.position[d] - A * D
        return X

    def update_position(self, X):
        self.position = (X[0] + X[1] + X[2]) / 3.0
        self.check_lim()
        self.update_fitness()

    def check_lim(self):
        for d in range(self.dim):
            if self.position[d] < self.lim[0]:
                self.position[d] = self.lim[0]
            elif self.position[d] > self.lim[1]:
                self.position[d] = self.lim[1]

    def update_fitness(self):
        self.fitness = self.obj_function(self.position, self.dim)


"""
    GreyWolfOptimization does not have any input control parameters, only the standart parameters in every swarm 
    algorithm
    Wolves are divided into alpha beta delta and omega -> three best solutions ( first best, second best and third best)
    are alpha beta and delta wolves, other solutions ( wolves ) are omega wolves
    Omega wolves update their position according to three best solutions ( alpha beta and delta)
    We sort the population and first three are three best, for every other wolf we update position
    Parameter a is control parameter, which is decreased in each iteration 
    Evaluate method -> first we sort the wolves then update global best solution after that we cycle for iter_limit
    iterations. For each iteration, first we set three best wolves, and parameter a, then we cycle through all wolves, 
    updating the positions of every omega. After update of every omega, we reset the flags on three best wolves and sort
    wolves population again. 
    Finally in each iteration we update global optima and best solution history
"""


class GreyWolfOptimization:
    def __init__(self, dimension, pop_size, iter_limit, obj_function, budget, limits):
        self.dimension = dimension
        self.limits = limits
        self.obj_function = obj_function
        self.pop_size = pop_size
        self.iter_limit = iter_limit
        self.grey_wolves = self.create_wolves()
        self.gbest = np.inf
        self.gbest_p = np.zeros(self.dimension)
        self.gbest_history = np.zeros([self.iter_limit, 1])
        self.gbest_history_p = np.zeros([self.iter_limit, self.dimension])
        self.alpha_wolf = None
        self.beta_wolf = None
        self.delta_wolf = None
        self.a = None
        self.budget = budget
        self.evaluations = 0
        self.iter_eval = [[-1, self.evaluations]]

    def create_wolves(self):
        wolves = [GreyWolf(self.limits, self.obj_function, self.dimension) for _ in range(self.pop_size)]
        return wolves

    def update_optimum(self):
        for wolf in self.grey_wolves:
            best_p = wolf.position
            best_f = wolf.fitness
            if best_f < self.gbest:
                self.gbest_p = deepcopy(best_p)
                self.gbest = deepcopy(best_f)

    def sort_wolves(self):
        self.grey_wolves.sort(key=lambda wolf: wolf.fitness)

    def set_alpha_beta_gamma_wolves(self):
        self.grey_wolves[0].set_alpha_status()
        self.grey_wolves[1].set_beta_status()
        self.grey_wolves[2].set_delta_status()
        self.alpha_wolf = self.grey_wolves[0]
        self.beta_wolf = self.grey_wolves[1]
        self.delta_wolf = self.grey_wolves[2]

    def reset_alpha_beta_gamma_wolves(self):
        self.grey_wolves[0].reset_alpha_status()
        self.grey_wolves[1].reset_beta_status()
        self.grey_wolves[2].reset_delta_status()

    def reduce_a(self, iteration):
        self.a = 2 - iteration * (2 / self.iter_limit)

    def evaluate(self):
        self.sort_wolves()
        self.update_optimum()
        for iteration in range(self.iter_limit):
            if self.evaluations < self.budget:
                self.set_alpha_beta_gamma_wolves()
                self.reduce_a(iteration)
                for grey_wolf in self.grey_wolves:
                    if self.evaluations > self.budget:
                        break
                    if grey_wolf.is_alpha or grey_wolf.is_beta or grey_wolf.is_delta:
                        continue
                    else:
                        grey_wolf.omega_wolf_move(self.alpha_wolf, self.beta_wolf, self.delta_wolf, self.a)
                        self.evaluations += 1
                self.reset_alpha_beta_gamma_wolves()
                self.sort_wolves()
                self.update_optimum_history(iteration)
                self.update_optimum()
                self.iter_eval.append([iteration, self.evaluations])
                # self.update_optimum_history(iteration)
            else:
                for l_iteration in range(self.iter_limit - 1, iteration - 1, -1):
                    self.gbest_history[l_iteration] = -105
                    self.gbest_history_p[l_iteration] = np.ones(self.dimension) * -105
                break

    def update_optimum_history(self, iteration):
        self.gbest_history[iteration] = self.gbest
        self.gbest_history_p[iteration] = self.gbest_p

    def print_results(self):
        print(f'Best fitness {self.gbest} at position {self.gbest_p}')
        print(f'History of global optima {self.gbest_history}')






