import numpy as np
import benchmark_functions
from copy import deepcopy
"""
    Antlion builds a trap at a random place, it waits until ant falls into its trap, ants are moving randomly around
    the search space
    Each antlion carries information like problem limits and dimension and problem objective funcion
    Antlion updates its position to ants position if the ant was caught and builds new trap at that location
"""


class AntLion:
    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)

    def update_position(self, position):
        for d in range(self.dim):
            self.position[d] = position[d]

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

    def build_new_trap(self, ant):
        if ant.fitness < self.fitness:
            self.update_position(ant.position)
            self.update_fitness(ant.fitness)


"""
    Ants are moving randomly through the search space, each ant carries information about problem with it, like limits,
    dimension and objective function, in addition it carries maximum iteration limit as well, because random walk of ant
    is generated for each iteration according to equation defined by author of algorithm
    Move method gets two arguments, first are two antlions and second is current iteration, first of two antlions is 
    the elite antlion - the one with the best fitness, second antlion is selected by roulette wheel selection randomly
    Ant randomly walks into trap set by antlion selected by roulette wheel, this is implemented in method random_walk
    Random walk method first calculates I, which is a ratio defined by equations by author of algorithm, this vaule
    is representing slide coefficient ( because ants are sliding down into antlion trap, which is shape of a cone )
    its value depends of iteration limit  and current iteration
    After I is computed, limits are set to represent convergence of ant towards antlion, thus these limits are then used
    for creating antlions trap ( creating boudaries around antlion ), then random walk is generated for every iteration
    and normalized with respect to boundaries, this operation is accomplished by equation defined by author of algorithm
    Elite antlion is influencing movement of all ants, therefore random walk component is generated for him as well
    Each random walk is composed of maximum iteration limit number of walks, ant selects for movement only those 
    components of these walks, which are on index of current iteration, with these values, ant updated its position this
    operation is also following equation, defined by author, which is implemented in update_position method
    After position update limits are checked and fitness of ant is updated
"""


class Ant:
    def __init__(self, limits, obj_function, dim, iter_limit):
        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.iter_limit = iter_limit

    def update_position(self, X_elite, X_random, iteration):
        r_part = np.zeros(self.dim)
        i = 0
        for walk in X_random:
            r_part[i] = walk[iteration]
            i += 1
        e_part = np.zeros(self.dim)
        i = 0
        for walk in X_elite:
            e_part[i] = walk[iteration]
            i += 1

        new_pos = (r_part + e_part) / 2
        for d in range(self.dim):
            self.position[d] = new_pos[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 set_I(self, iteration):
        I = 1
        if iteration > self.iter_limit / 10:
            I = 1 + 100 * (iteration / self.iter_limit)
        if iteration > self.iter_limit / 2:
            I = 1 + 1000 * (iteration / self.iter_limit)
        if iteration > self.iter_limit * (3 / 4):
            I = 1 + 10000 * (iteration / self.iter_limit)
        if iteration > self.iter_limit * 0.9:
            I = 1 + 100000 * (iteration / self.iter_limit)
        if iteration > self.iter_limit * 0.95:
            I = 1 + 1000000 * (iteration / self.iter_limit)
        return I

    def update_limit(self, limit, position):
        n_lim = np.zeros([self.dim, 1])
        for d in range(self.dim):
            n_lim[d] = limit + position[d]
        return n_lim

    def random_walk(self, antlion, iteration):
        I = self.set_I(iteration)
        lower_bound = self.lim[0] / I
        upper_bound = self.lim[1] / I
        if np.random.rand() < 0.5:
            lower_bound = self.update_limit(lower_bound, antlion.position)
        else:
            lower_bound = self.update_limit(-lower_bound, antlion.position)
        if np.random.rand() >= 0.5:
            upper_bound = self.update_limit(upper_bound, antlion.position)
        else:
            upper_bound = self.update_limit(-upper_bound, antlion.position)
        X = []
        for d in range(self.dim):
            r = np.random.rand(self.iter_limit, 1)
            r = r > 0.5
            r_walk = np.cumsum(2 * r - 1)
            a = np.min(r_walk)
            b = np.max(r_walk)
            c = lower_bound[d]
            d = upper_bound[d]
            r_walk_norm = ((r_walk - a) * (d - c)) / (b - a) + c
            X.append(r_walk_norm)
        return X

    def move(self, antlions, iteration):
        r_w_around_elite = self.random_walk(antlions[0], iteration)
        r_w_around_random = self.random_walk(antlions[1], iteration)
        self.update_position(r_w_around_elite, r_w_around_random, iteration)

"""
    AntLionOptimization is utilizing hunting mechanism of antlions, it uses two classes, ants and antlions
    First population of ants and antlions are generated at random positions
    Then antlions are sorted and the best is selected as elite and global best solution. After this operation we loop
    for iter_limit ( maximum iteration number - budget ) times through all ants and update their position with elite
    antlion and antlion selected by roulette wheel ( implemented in roulette_wheel method )
    After update of position of every ant comes the last stage of hunt, in which antlions update their positions with
    respect to ant -> that means if ant becomes fittier than antlion, it is persumed it was consumed by antlion, which
    updates its position to position of that ant and builds a new trap
    We update global optima, elite and history and repeat for another iteration 
"""


class AntLionOptimization:
    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.ants = self.create_ants()
        self.antlions_swarm = self.create_antlions()
        self.elite = None
        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.budget = budget
        self.evaluations = 0
        self.iter_eval = [[-1, self.evaluations]]

    def create_ants(self):
        ants = [Ant(self.limits, self.obj_function, self.dimension, self.iter_limit) for _ in range(self.pop_size)]
        return ants

    def set_elite(self):
        self.elite = self.antlions_swarm[0]

    def create_antlions(self):
        antlion = [AntLion(self.limits, self.obj_function, self.dimension) for _ in range(self.pop_size)]
        return antlion

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

    def update_antlions(self, swarm):
        for a_index in range(self.pop_size):
            self.antlions_swarm[a_index].build_new_trap(swarm[a_index])

    def sort(self, swarm):
        swarm.sort(key=lambda insect: insect.fitness)

    def roulette_wheel(self):
        antlion_fitnessess = [(1 / antlion.fitness) for antlion in self.antlions_swarm]
        cumsum = np.cumsum(antlion_fitnessess)
        p = np.random.rand() * cumsum[-1]
        for index in range(len(cumsum)):
            if cumsum[index] > p:
                return index
        return 0

    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 print_partial(self, iteration):
        if (iteration % 100 == 0):
            print(f'Fitness {self.gbest} at iteration {iteration}')

    def evaluate(self):
        self.sort(self.antlions_swarm)
        self.update_optimum()
        self.set_elite()
        for iteration in range(self.iter_limit):
            if self.evaluations < self.budget:
                self.update_optimum_history(iteration)
                for ant in self.ants:
                    if self.evaluations > self.budget:
                        break
                    ant.move([self.elite, self.antlions_swarm[self.roulette_wheel()]], iteration)
                    self.evaluations += 1
                concatenated_swarm = deepcopy(self.antlions_swarm + self.ants)
                self.sort(concatenated_swarm)
                self.update_antlions(concatenated_swarm)
                self.update_optimum()
                self.set_elite()
                self.iter_eval.append([iteration, self.evaluations])
                # self.print_partial(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

