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

"""
    Dragonfly is generated at random position from which is counted fitness, each dragonfly has neighbours, and several
    parameters defined by author, which controls the algorithm and vector deltaX, which is step vector, carrying numbers
    defining movement in each dimension
    Neighbours are discovered in find_neighbours method, dragonfly searches in its close proximity, which is defined by
    parameter r, which is computed in DragonflyOptimization, and changes in every iteration
    Control parameters, S, A, C, F and E are computed in method compute_principles, they are defined by equations 
    defined by author, dragonfly also computes distance to food source ( which is the best solution - dragonfly ) and
    enemy ( which is the worst solution - dragonfly ), these parameters are then used in position update - movement
    Dragonfly can move according to three principles, every is defined by separate method, which one it selects depends
    on number of neighbours and by parameter r
    Each movement method is implemented through respective equations, defined by author
    In every movement method, not only position is updated, but deltaX vector is updated too
    To make sure that movement and respective steps from deltaX vector do not go out of bounds, limits are checked 
"""


class Dragonfly:
    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.neighbours = []
        self.deltaX = np.random.uniform(self.lim[0], self.lim[1], self.dim)
        self.S = 0
        self.A = 0
        self.C = 0
        self.F = 0
        self.E = 0
        self.distance_to_food = 0
        self.distance_to_enemy = 0

    def compute_distance(self, food_or_enemy):
        return np.abs(food_or_enemy - self.position)

    def find_neighbours(self, dragonflies, r):
        for neighbour in dragonflies:
            distance = np.abs(self.position - neighbour.position)
            if np.all(distance <= r) and np.all(distance != 0):
                self.neighbours.append(neighbour)

    def compute_principles(self, food, enemy, r):
        if len(self.neighbours) > 1:
            self.S = np.sum([neighbour.position - self.position for neighbour in self.neighbours], axis=0)
            self.S = -self.S
            self.A = np.sum([neighbour.deltaX for neighbour in self.neighbours], axis=0) / len(self.neighbours)
            self.C = np.sum([neighbour.deltaX for neighbour in self.neighbours], axis=0) - self.position
        else:
            self.S = np.zeros(self.dim)
            self.A = self.deltaX
            self.C = self.position - self.position
        self.distance_to_food = self.compute_distance(food.position)
        if np.all(self.distance_to_food <= r):
            self.F = food.position - self.position
        else:
            self.F = np.zeros(self.dim)
        self.distance_to_enemy = self.compute_distance(enemy.position)
        if np.all(self.distance_to_enemy <= r):
            self.E = enemy.position + self.position
        else:
            self.E = np.zeros(self.dim)

    def move_levy(self):
        levy_lambda = 3 / 2
        sigma = (gamma(1 + levy_lambda) * np.sin(np.pi * levy_lambda / 2) / (
                    gamma((1 + levy_lambda) / 2) * levy_lambda * 2 ** ((levy_lambda - 1) / 2))) ** (1 / levy_lambda)
        u = np.random.randn(self.dim) * sigma
        v = np.random.randn(self.dim)
        step = u / np.abs(v) ** (1 / levy_lambda)
        self.position = 0.01 * step * self.position
        self.check_lim()
        self.update_fitness()
        self.deltaX = np.zeros(self.dim)

    def move_var_one(self, w, delta_m):
        movement = w * self.deltaX + np.random.uniform(0, 1, self.dim) * self.A + \
                           np.random.uniform(0, 1, self.dim) * self.C + np.random.uniform(0, 1, self.dim) * self.S
        movement = np.clip(movement, -1 * delta_m, delta_m)
        self.position = self.position + movement
        self.check_lim()
        self.update_fitness()
        self.deltaX = movement

    def move_var_two(self, a, c, s, f, e, w, delta_m):
        movement = (a * self.A + c * self.C + s * self.S + f * self.F + e * self.E) + w * self.deltaX
        movement = np.clip(movement, -1 * delta_m, delta_m)
        self.position = self.position + movement
        self.check_lim()
        self.update_fitness()
        self.deltaX = movement

    def move(self, w, a, c, s, f, e, r, delta_m):
        if np.any(self.distance_to_food > r):
            if len(self.neighbours) > 1:
                self.move_var_one(w, delta_m)
            else:
                self.move_levy()
        else:
            self.move_var_two(a, c, s, f, e, w, delta_m)
        self.clear_neighbours()

    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 clear_neighbours(self):
        self.neighbours = []


"""
    Dragonfly algorithm creates dragonfly swarm at random positions
    Algorithm parameters s a c f e w and cc are computed in compute parameters method, by equations defined by author,
    these are weigths which drangonfly uses to compute their 
    movement
    Parameter r is defining surroundings of dragonfly, it uses it to find its neighbours and to determine method of 
    movement
    In each iteration, there is enemy and food source (these are the worst and best dragonfly in the swarm dragonfly tries
    to stay away from enemy and get close to food source) and the parameters mentioned before are also recomputed again
    In each iteration, we loop through dragonfly, make it find its neighbours, compute primitive principles which will
    be vital in movement equations and finally, we make dragonfly update its position, by which method it selects 
    according to parameters passed
    After iteration we sort dragonflies, update optimum and history and repeat the cycle
"""


class DragonflyAlgorithm:
    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.dragonfly_swarm = self.create_dragonflies()
        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.r = 0
        self.d_max = (self.limits[1] - self.limits[0]) / 10
        self.enemy = None
        self.food_source = None
        self.s = 0
        self.a = 0
        self.c = 0
        self.f = 0
        self.e = 0
        self.w = 0
        self.cc = 0
        self.budget = budget
        self.evaluations = 0
        self.iter_eval = [[-1, self.evaluations]]

    def create_dragonflies(self):
        dragonfiles = [Dragonfly(self.limits, self.obj_function, self.dimension) for _ in range(self.pop_size)]
        return dragonfiles

    def update_optimum(self):
        for dragonfly in self.dragonfly_swarm:
            best_p = dragonfly.position
            best_f = dragonfly.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 select_predator_and_food(self):
        self.enemy = deepcopy(self.dragonfly_swarm[-1])
        self.food_source = deepcopy(self.dragonfly_swarm[0])

    def sort(self):
        self.dragonfly_swarm.sort(key=lambda df: df.fitness)

    def compute_parameters(self, iteration):
        self.r = (self.limits[1] - self.limits[0]) / 4 + ((self.limits[1] - self.limits[0])
                                                          * ((iteration + 1) / self.iter_limit) * 2)
        self.w = 0.9 - iteration * ((0.9 - 0.4) / self.iter_limit)
        self.cc = 0.1 - (iteration + 1) * ((0.1 - 0) / self.iter_limit / 2)
        self.cc = 0 if self.cc < 0 else self.cc
        self.s = 2 * np.random.rand() * self.cc
        self.a = 2 * np.random.rand() * self.cc
        self.c = 2 * np.random.rand() * self.cc
        self.f = 2 * np.random.rand()
        self.e = deepcopy(self.cc)

    def evaluate(self):
        self.sort()
        self.update_optimum()
        for iteration in range(self.iter_limit):
            if self.evaluations < self.budget:
                self.update_optimum_history(iteration)
                self.select_predator_and_food()
                self.compute_parameters(iteration)
                for dragonfly in self.dragonfly_swarm:
                    if self.evaluations > self.budget:
                        break
                    dragonfly.find_neighbours(self.dragonfly_swarm, self.r)
                    dragonfly.compute_principles(self.food_source, self.enemy, self.r)
                    dragonfly.move(self.w, self.a, self.c, self.s, self.f, self.e, self.r, self.d_max)
                    self.evaluations += 1
                self.sort()
                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

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

