import numpy as np
import benchmark_functions
from copy import deepcopy

"""
    Each firefly has randomness - alpha, atractivness beta and absorbtion coefficent gamma
    Posision of firefly is generated randomly, from position intensity is computed - which is in fact its fitness
    Firefly starts with 0 distance to better -> this distance is distance to firefly which has better intensity
    than firefly, it changes throughout the iterations as the firefly constantly changes it's position towards
    the better fireflies
    Apart from methods that introduce computing of epsilon and beta, we also implemented update_intensity method, which
    is called after firefly changed its position in FireflyAlgorithm.evaluate() method
    Optional method Reduce randomness is introduced, it will decrease value of alpha according to two possible principles
    first is by utilizing theta second by static value, this method is not used, because from certain number of runs
    I came to conclusion that algorithm performs better with static value of alpha
"""


class Firefly:
    def __init__(self, alpha, beta, gamma, limits, obj_function, dim):
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        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.intensity = 1 / self.fitness
        self.distance_to_better = 0

    def compute_distance(self, n_position):
        self.distance_to_better = np.linalg.norm(self.position - n_position)

    def move_to_better(self, n_position):
        for d in range(self.dim):
            self.position[d] = self.position[d] + self.compute_beta() *\
                               (n_position[d] - self.position[d]) + self.alpha * self.compute_epsilon()
        self.check_lim()

    def compute_epsilon(self):
        return np.random.uniform(0, 1) - 0.5

    def compute_beta(self):
        return np.exp(-self.gamma * (self.distance_to_better ** 2))

    def update_intensity(self):
        self.fitness = self.obj_function(self.position, self.dim)
        self.intensity = 1 / self.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 reduce_randomness(self, max_iter):
    #     theta = 10 ** (-5/max_iter)
    #     self.alpha = self.alpha * theta
    #     # or
    #     # self.alpha = self.alpha * 0.97

    def update_alpha(self, alpha):
        self.alpha = alpha


"""
    Firefly algorithm initializes firefly swarm in its constructor, passing them all the values necessary 
    It includes best global solution and best solutions history and their positions as well
    After initialization of fireflies Evaluate method is called, where firefly tries to change its position according
    to intensity of light emmited by other fireflies, we make copy of swarm of fireflies and remove the one currently
    searching neighbours for position update
    After all fireflies updated their positions, we update the global best solution so far by utiliznig greedy selection
    and global best solutions history as well
"""


class FireflyAlgorithm:
    def __init__(self, dim, pop_size, iter_limit, obj_func, budget, limits, alpha, beta, gamma):
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        self.limits = limits
        self.obj_function = obj_func
        self.dim = dim
        self.iter_limit = iter_limit
        self.pop_size = pop_size
        self.firefly_swarm = self.init_firefly()
        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 update_optimum(self):
        for firefly in self.firefly_swarm:
            best_p = firefly.position
            best_f = firefly.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 init_firefly(self):
        fireflies = [Firefly(self.alpha, self.beta, self.gamma,  self.limits, self.obj_function, self.dim)
                 for _ in range(self.pop_size)]
        return fireflies

    def reduce_randomness(self, max_iter):
        theta = 10 ** (-5/max_iter)
        self.alpha = self.alpha * theta
        # or
        # self.alpha = self.alpha * 0.97
        for firefly in self.firefly_swarm:
            firefly.update_alpha(self.alpha)

    def evaluate(self):
        self.update_optimum()
        for iteration in range(self.iter_limit):
            self.update_optimum_history(iteration)
            if self.evaluations < self.budget:
                for firefly in self.firefly_swarm:
                    if self.evaluations > self.budget:
                        break
                    neighbour_firefly_swarm = self.firefly_swarm.copy()
                    neighbour_firefly_swarm.remove(firefly)
                    for neighbour_firefly in neighbour_firefly_swarm:
                        if neighbour_firefly.intensity > firefly.intensity:
                            firefly.compute_distance(neighbour_firefly.position)
                            firefly.move_to_better(neighbour_firefly.position)
                            firefly.update_intensity()
                            self.evaluations += 1
                self.reduce_randomness(self.iter_limit)
                # 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.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}')

