import numpy as np
import benchmark_functions
from copy import deepcopy

"""
    Class Whale is defined by its position and fitness, position is randomly generated, fitness computed by objective
    function with position as argument, each Whale carries limits of optimization problem and dimension of problem with
    it, finally it has a value D which is vital in position updating and value b which is a constant defining shape of
    logaritmic spiral used in bubble net attack
    Encircle prey: this method is implementing first type of position update, which is working on principle, that the 
    new position of whale is computed by using the best whale in swarm and arguments A and D which are passed to it, 
    D argument is computed by update_D and uses best whale to get its value => in short, this method means that whale
    is trying to improve its position by following the best whale in swarm
    Search for prey: this method is very simillar to encircle prey method, however instead of best whale, it uses 
    random selected whale from whale swarm, other than that, it works exactly the same => in short, whales are still
    looking for food
    Bubble net attack: this method is modelling the principle of hunting method of humpack whales, the bubble net attack
    D argument is computed differently than in previous two methods, it uses current position of whale and best whale, 
    new position is then acquired by mathematical equation defined by inventor of algorithm, using the best whale 
    position and several other arguments -> b, which was described earlier, and l which is a random number in interval
    [-1,1] 
    After each position update, any method for update is used, limits are checked and new fitness is computed
"""


class Whale:
    def __init__(self, limits, obj_function, dim, b):
        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.D = np.zeros(self.dim)
        self.b = b

    def encircle_prey(self, A, C, best_whale):
        self.update_D(C, best_whale, False)
        for d in range(self.dim):
            self.position[d] = best_whale.position[d] - A * self.D[d]
        self.check_lim()
        self.update_fitness()

    def search_for_prey(self, A, C, whales):
        whales.remove(self)
        whale = np.random.choice(whales)
        self.update_D(C, whale, False)
        for d in range(self.dim):
            self.position[d] = whale.position[d] - A * self.D[d]
        self.check_lim()
        self.update_fitness()

    def bubble_net_attack(self, C, best_whale):
        self.update_D(C, best_whale, True)
        l = np.random.uniform(-1.0, 1.0, size=self.dim)
        for d in range(self.dim):
            self.position[d] = self.D[d] * np.exp(self.b * l[d]) * np.cos(2.0 * np.pi * l[d]) + best_whale.position[d]
        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)

    def update_D(self, C, whale, attack):
        if attack:
            for d in range(self.dim):
                self.D[d] = np.abs(whale.position[d] - self.position[d])
        else:
            for d in range(self.dim):
                self.D[d] = np.abs(C * whale.position[d] - self.position[d])

"""
    Whale Optimization carries arguments describing problem like dimension and limits of problem, its objective function
    and others like population size and iteration limit which is in fact budget. There is another parameter, b which is
    specific for WOA and is defining shape of spiral used in bubble net attack
    First whales are created, at random positions, then we sort them and update global optima with values of the best 
    whale
    then in each iteration, we loop through all whales and compute several parameters
    Parameter A in update_A, by defined equation
    Parameter C in update_C, by defined equation
    Parameter a in reduce_a, which is decreasing its value in every iteration from 2 to 0
    Parameter p in update_p, which is a random number used to determine the method of position of whale update
    Evaluate method: for each whale in iteration, we compute parameters, use one of defined methods to update its 
    position and after looping through every whale, we sort all whales in swarm and update global optima
"""


class WhaleOptimization:
    def __init__(self, dimension, pop_size, iter_limit, obj_function, budget, limits, b):
        self.b = b
        self.dimension = dimension
        self.limits = limits
        self.obj_function = obj_function
        self.pop_size = pop_size
        self.iter_limit = iter_limit
        self.whale_swarm = self.create_whales()
        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.a = None
        self.C = None
        self.p = None
        self.A = None
        self.budget = budget
        self.evaluations = 0
        self.iter_eval = [[-1, self.evaluations]]

    def create_whales(self):
        whales = [Whale(self.limits, self.obj_function, self.dimension, self.b) for _ in range(self.pop_size)]
        return whales

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

    def sort_whales(self):
        self.whale_swarm.sort(key=lambda whale: whale.fitness)

    def update_A(self):
        r = np.random.random()
        self.A = 2 * self.a * r - self.a

    def update_C(self):
        c = np.random.random()
        self.C = 2 * c

    def update_p(self):
        self.p = np.random.random()

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

    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}')

    def evaluate(self):
        self.sort_whales()
        self.update_optimum()
        for iteration in range(self.iter_limit):
            if self.evaluations < self.budget:
                self.reduce_a(iteration)
                for whale in self.whale_swarm:
                    if self.evaluations > self.budget:
                        break
                    self.update_A()
                    self.update_C()
                    self.update_p()
                    if self.p > 0.5:
                        if np.abs(self.A) < 1:
                            best_whale = self.whale_swarm[0]
                            whale.encircle_prey(self.A, self.C, best_whale)
                            self.evaluations += 1
                        else:
                            whale.search_for_prey(self.A, self.C, self.whale_swarm.copy())
                            self.evaluations += 1
                    else:
                        best_whale = self.whale_swarm[0]
                        whale.bubble_net_attack(self.C, best_whale)
                        self.evaluations += 1
                self.sort_whales()
                self.update_optimum_history(iteration)
                self.update_optimum()
                self.iter_eval.append([iteration, self.evaluations])
            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

