import benchmark_functions
import numpy as np
from copy import deepcopy

"""
    Each bat has several control parameters, those are its loudnes and pulse rate, which are first randomly generated 
    and then updated in iteration if cetrain condition is met ( that is if the new solution position fitness is better
    than previous one, and loudness is smaller than randomly generated number -> equation defined by author
    Loudness and pulse rate are in its respecitve dimensions, which are set in algorithm initialization and thus serve
    as static control parameters, similiar to these two, each bat has a frequency at which it emmits pulses, this is
    computed by its respecitive equation defined by author, while remaining in defined limits
    Alpha and gamma are parameters which also serve for updating of loudness and pulse rate, making them another two
    static parameters of algorithm
    Bat tries to update its position, first frequency is computed, from it velocity is computed and by velocity, 
    position is computed. Then position can be furhter randomized by random walk, if random number is greater than 
    pulse emision rate. Bat updates its position only if the fitness of new position is better than old one, and loudness
    is smaller than random generated number. Also, if this condition is met, not only position and fitness are updated,
    parameters -> loudness and pulse rate are updated as well, with utilization of alpha and gamma
"""


class Bat:
    def __init__(self, alpha, gamma, Alim, rlim, fqlim, limits, dim, obj_function):
        self.alpha = alpha
        self.gamma = gamma
        self.Alim = Alim
        self.rlim = rlim
        self.fqlim = fqlim
        self.limits = limits
        self.dim = dim
        self.obj_function = obj_function
        self.position = np.random.uniform(self.limits[0], self.limits[1], self.dim)
        self.velocity = np.random.uniform(self.limits[0], self.limits[1], self.dim)
        self.fitness = self.obj_function(self.position, self.dim)
        self.loudness = np.random.uniform(self.Alim[0], self.Alim[1])
        self.pulse_rate = np.random.uniform(self.rlim[0], self.rlim[1])
        self.frequency = 0
        self.potential_position = np.zeros(self.dim)
        self.potential_fitness = 0

    def update_loudness(self):
        self.loudness = self.alpha * self.loudness

    def update_pulse_rate(self, iteration):
        self.pulse_rate = self.pulse_rate * (1 - np.exp(-self.gamma * (iteration + 1)))

    def update_frequency(self):
        self.frequency = self.fqlim[0] + (self.fqlim[1] - self.fqlim[0]) * np.random.uniform()

    def update_velocity(self, best_bat):
        self.velocity = self.velocity + (self.position - best_bat.position) * self.frequency

    def update_position(self):
        self.potential_position = self.position + self.velocity
        self.check_lim()
        self.update_fitness()

    def move(self, best_bat, mean_loudness, iteration):
        evaluation = 0
        self.update_frequency()
        self.update_velocity(best_bat)
        self.update_position()
        evaluation += 1
        if np.random.uniform() > self.pulse_rate:
            self.potential_position = self.potential_position + mean_loudness * np.random.normal(-1, 1)
            self.check_lim()
            self.update_fitness()
            evaluation += 1
        if self.potential_fitness < self.fitness and np.random.rand() < self.loudness:
            self.position = deepcopy(self.potential_position)
            self.fitness = deepcopy(self.potential_fitness)
            self.update_parameters(iteration)
        return evaluation

    def update_parameters(self, iteration):
        self.update_loudness()
        self.update_pulse_rate(iteration)

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

    def update_fitness(self):
        self.potential_fitness = self.obj_function(self.potential_position, self.dim)


"""
    Bat algorithm has several static control parameters -> those are alpha gamma and parameter limits-> limits of
    loudness, pulse rate and frequency. 
    In each optimum update, we also update best bat, which is used to compute new position of bats
    Evaluate: after sorting all bats and updating best bat, we loop for iter_limit times through all bats, updating 
    their positions. After every iteration, swarm is sorted, with the best individual at the top, optimum is updated and
    history is updated, best bat as well
"""


class BatAlgorithm:
    def __init__(self, dimension, pop_size, iter_limit, obj_function, budget, limits, gamma, alpha, param_limits):
        self.gamma = gamma
        self.alpha = alpha
        self.param_limits = param_limits
        self.dimension = dimension
        self.limits = limits
        self.obj_function = obj_function
        self.pop_size = pop_size
        self.iter_limit = iter_limit
        self.bat_swarm = self.create_bats()
        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.best_bat = None
        self.mean_loudness = 0
        self.budget = budget
        self.evaluations = 0
        self.iter_eval = [[-1, self.evaluations]]

    def create_bats(self):
        bats = [Bat(self.alpha, self.gamma, self.param_limits[0], self.param_limits[1], self.param_limits[2],
                    self.limits, self.dimension, self.obj_function) for _ in range(self.pop_size)]
        return bats

    def update_optimum(self):
        for bat in self.bat_swarm:
            best_p = bat.position
            best_f = bat.fitness
            if best_f < self.gbest:
                self.best_bat = deepcopy(bat)
                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 sort(self):
        self.bat_swarm.sort(key=lambda bat: bat.fitness)

    def update_mean_loudness(self):
        self.mean_loudness = np.mean([bat.loudness for bat in self.bat_swarm])

    def evaluate(self):
        self.sort()
        self.update_optimum()

        for iteration in range(self.iter_limit):
            if self.evaluations < self.budget:
                self.update_mean_loudness()
                self.update_optimum_history(iteration)
                for bat in self.bat_swarm:
                    if self.evaluations > self.budget:
                        break
                    else:
                        self.evaluations += bat.move(self.best_bat, self.mean_loudness, iteration)
                self.sort()
                self.update_optimum()
                self.iter_eval.append([iteration, self.evaluations])
                # self.update_mean_loudness()
                # self.update_optimum_history(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

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

