import numpy as np
import benchmark_functions
import math
from copy import deepcopy

"""
    Cuckoo has values of alpha, which is used for sacling of problem, step which is value computed from levy flight 
    random move and levy lambda, which is used in levy flight when computing step value
    Fitness of cuckoo is determined from its position
    Cuckoo, nest and egg are the same entity, we do not diferentiate between them
    Move: in this method cuckoo performes random move -> to update position of cuckoo with usade of levy flight, 
    it uses the step value which is computed in compute_step_size called at the beggining of this method, after that
    limits of new position are checked and its fitness is updated
    Check_lim: checks the limits for each dimension in position
    Compute_step_size: computation of levy component, u and v are obtained through normal distribution, sigma_u is 
    standart deviation defined through its respective equation, step size is computed also accoridng to its defined
    equation, sigma_v is always one
    Create_new_nest is called when egg in host nest is discovered, thus cuckoo needs to lay a new one in another nest
    this is simulated through creating another nest at random position, in which cuckoo lays its egg ( again we do not
    make diferences between egg and cuckoo and nest)
"""


class Cuckoo:
    def __init__(self, alpha, levy_lambda, limits, dim, objfunc):
        self.alpha = alpha
        self.levy_lambda = levy_lambda
        self.step = 0
        self.lim = limits
        self.dim = dim
        self.obj_func = objfunc
        self.position = np.random.uniform(self.lim[0], self.lim[1], self.dim)
        self.fitness = self.obj_func(self.position, self.dim)

    def move(self, best, iteration):
        self.compute_step_size()
        potential_position = np.zeros(self.dim)
        for d in range(self.dim):
            potential_position[d] = self.position[d] + self.alpha / np.sqrt(iteration + 1) * np.sign(
                np.random.random() - 0.5) * self.step[d] * (self.position[d] - best[d])
        self.check_lim(potential_position)
        potential_fitness = self.obj_func(potential_position, self.dim)

        if potential_fitness < self.fitness:
            self.update_position(potential_position)
            self.update_fitness(potential_fitness)

    def update_fitness(self, potential_fitness):
        self.fitness = potential_fitness

    def update_position(self, potential_position):
        self.position = potential_position

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

    def compute_step_size(self):
        sigma_v = 1
        sigma_u = np.power((math.gamma(1 + self.levy_lambda) * np.sin((np.pi * self.levy_lambda) / 2))
                           / math.gamma((1 + self.levy_lambda) / 2) * np.power(2, (self.levy_lambda - 1) / 2),
                           1 / self.levy_lambda)
        u = np.random.normal(0, sigma_u, size=self.dim)
        v = np.random.normal(0, sigma_v, size=self.dim)
        self.step = u / abs(v) ** (1 / self.levy_lambda)

    def create_new_nest(self):
        self.position = np.random.uniform(self.lim[0], self.lim[1], self.dim)
        self.fitness = self.obj_func(self.position, self.dim)


"""
    CuckooSearch contains values necessary for algorithm control like alpha, levy_labmda and pa
    First two are described in Cuckoo object, pa is probability of host discovering cuckoo egg, this is described in
    method abandon_nest
    First cuckoo swarm is initiated, with its respective parameters
    Evaluate: this method is main optimization method calling respective methods of algorithm, first, optimum is updated
    because global best solution is needed in levy flight, then for each cuckoo, move is called
    Move is described in Cuckoo, it is performing random walk to new position according to global best position
    After change of position of every cuckoo, we again update optimum, to save the best one
    Abandon nest: this method simulates discovery of egg by host according to probability pa which is compared with 
    r which is number obtained by uniform distribution, if its discovered, new nest must be created by cuckoo,
    this is explained in Cuckoo object
    after this method, we again update optimum in hope of finding a better solution, becuse some eggs in nests were
    discovered, new nests were infested by cuckoos at new position, thus becoming new competitors for best solution
    finally we update history of global best solution in current iteration 
"""


class CuckooSearch:
    def __init__(self, dim, pop_size, iter_limit, obj_func, budget, limits, alpha, levy_lambda, pa):
        self.alpha = alpha
        self.levy_lambda = levy_lambda
        self.dim = dim
        self.limits = limits
        self.obj_func = obj_func
        self.iter_limit = iter_limit
        self.pop_size = pop_size
        self.cuckoo_swarm = self.init_cuckoos()
        self.pa = pa
        self.gbest = np.inf
        self.gbest_p = np.zeros(self.dim)
        self.gbest_history = np.zeros([self.iter_limit, 1])
        self.gbest_history_p = np.zeros([self.iter_limit, self.dim])
        self.budget = budget
        self.evaluations = 0
        self.iter_eval = [[-1, self.evaluations]]

    def init_cuckoos(self):
        cuckoos = [Cuckoo(self.alpha, self.levy_lambda, self.limits, self.dim, self.obj_func)
                   for _ in range(self.pop_size)]
        return cuckoos

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

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

    def abandon_nest(self):
        for cuckoo in self.cuckoo_swarm:
            r = np.random.rand()
            if r < self.pa:
                cuckoo.create_new_nest()
                self.evaluations += 1

    def evaluate(self):
        self.update_optimum()
        for iteration in range(self.iter_limit):
            if self.evaluations < self.budget:
                self.update_optimum_history(iteration)
                for cuckoo in self.cuckoo_swarm:
                    if self.evaluations > self.budget:
                        break
                    cuckoo.move(self.gbest_p, iteration)
                    self.evaluations += 1
                self.update_optimum()
                self.abandon_nest()
                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.dim) * -105
                break

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

