import numpy as np
import benchmark_functions
from copy import deepcopy

"""
    This bee is starting the ABC algorithm optimization phase,
    Position of this bee ( same as onlooker ) is randomly generated, from this position, fvalue (objective function value)
    and fitness are calculated (IMPORTANT NOTE: in ABC, fitness and objective function value are not the same values, fitness
    is calculated from fvalue (objective function value) in calculate_fitness method, trial is set to zero, it is incremented
    each time bee fails to find better solution ( position with better fitness than its actual ),
    mtrial is the maximum number of trials, when exceeded bee will enter scout phase for discard and reset
    explore method is generating new solutions based on current position and position of randomly selected agent, which
    must not be the same as the current position, therefore we create a copy and remove current position from EmployedBees
    after generating new solution it is compared with current values to determine if its better in p_and_f_update method
    check_lim checks if the new position is in the range of function limits, c_probability computes probability of bee
    according to the defined equation and p_and_f_update will update bee's position and fitness with values of the better
    solution generated by bee in explore method
"""


class EmployedBee:
    def __init__(self, dimension, lim, objfunc, mtrial):
        self.dimension = dimension
        self.lim = lim
        self.position = np.random.uniform(self.lim[0], self.lim[1], self.dimension)
        self.objfunc = objfunc
        self.fitness = self.objfunc(self.position, self.dimension)
        self.f = self.calculate_fitness()
        self.trial = 0
        self.prob = 0
        self.mtrial = mtrial

    def calculate_fitness(self):
        if self.fitness <= 0:
            return 1/(1 + self.fitness)
        else:
            return 1 + np.abs(self.fitness)

    def explore(self, bees):
        evaluations = 0
        if self.trial <= self.mtrial:
            c_bees = bees.copy()
            c_bees.remove(self)
            partner = np.random.choice(c_bees)
            phi = np.random.uniform(-1, 1, size=self.dimension)
            new_position = self.position + (self.position - partner.position) * phi
            new_position = self.check_lim(new_position)
            self.p_and_f_update(new_position)
            evaluations += 1
        return evaluations

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

    def p_and_f_update(self, new_pos):
        new_fitness = self.objfunc(new_pos, self.dimension)
        if self.fitness >= new_fitness:
            for d in range(self.dimension):
                self.position[d] = new_pos[d]
            self.fitness = new_fitness
            self.f = self.calculate_fitness()
            self.trial = 0
        else:
            self.trial += 1

    def c_probability(self, fitness_sum):
        self.prob = self.f / fitness_sum
        # alternatively when fitness_sum = min_fitness
        # self.prob = 0.9 * (fitness_sum / self.fitness) + 0.1


"""
    Onlooker bee shares similar values and methods with Employed bee, therefore they wont be described again
    Onlooker bee is starting to exploit after EmployedBees have generated new solutions
    It works on principle of probability -> to generate a new solution from EmployedBee, it checks if the probability
    counted from EmployedBee fitness, is higher than random number, if yes, it will try to generate new better solution
    if this solution is better in fitness than its current one, it will make it its new position and fitness
    Even if OnlookerBee has an probability, its purely to serve just for reset in main algorithm, OnlookerBee is not
    utilizing this value, its just to counter error
    even though exploit looks similar to explore -> exploit is trying to generate new solution only when condition is met
    it is important to note that OnlookerBee still uses position of EmployedBee to generate new solution, not its own
    BUT when the new solution is better (position and fittness better than current Onlooker Bee's), it updates its values
    with this new solution values
"""


class OnlookerBee:
    def __init__(self, dimension, lim, objfunc, mtrial):
        self.dimension = dimension
        self.lim = lim
        self.position = np.random.uniform(self.lim[0], self.lim[1], self.dimension)
        self.objfunc = objfunc
        self.fitness = self.objfunc(self.position, self.dimension)
        self.f = self.calculate_fitness()
        self.trial = 0
        self.mtrial = mtrial
        self.prob = 0

    def calculate_fitness(self):
        if self.fitness <= 0:
            return 1/(1 + self.fitness)
        else:
            return 1 + np.abs(self.fitness)

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

    def exploit(self, bees):
        evaluations = 0
        c_bees = bees.copy()
        for bee in c_bees:
            if np.random.uniform(0, 1) < bee.prob:
                if self.trial <= self.mtrial:
                    c_bees.remove(bee)
                    partner = np.random.choice(c_bees)
                    phi = np.random.uniform(-1, 1, size=self.dimension)
                    new_pos = bee.position + (bee.position - partner.position) * phi
                    new_pos = self.check_lim(new_pos)
                    self.p_and_f_update(new_pos)
                    evaluations += 1
        return evaluations

    def p_and_f_update(self, new_pos):
        new_fitness = self.objfunc(new_pos, self.dimension)
        if self.fitness >= new_fitness:
            for d in range(self.dimension):
                self.position[d] = new_pos[d]
            self.fitness = new_fitness
            self.calculate_fitness()
            self.trial = 0
        else:
            self.trial += 1


"""
    Colony composed of two bee types - onlookers and employed, scouts are not represented, scout phase is excecuted
    through ArtificalBeeColony class: trial_limit - how many times will bee try to find better solution, if it exceeds
    the trial limit, it will no longer participate in search for better solution --> in exploit and explore methods
    Two bee types created with respective methods -> their positions are randomly generated
    Evaluate has several phases, first is EmployedBee phase, which is described in its class, second is optimum update
    which will save the best so far solution from ebees
    Compute probability is vital for onlooker phase, because onlookers decide which solution( EmployedBee) to explore
    according to probablity which is counted from fitness, after onlooker phase, we will again try to update optimum
    this time using solutions from obees, after this phase we will save the global best solution and update history
    before moving to the scout phase, it is necessary to do this, because scout phase can drop even the best solution
    Scout phase will drop bees which exceeded trial_limit, it will generate new bee at their place, with random position
    and will count their fitness
"""


class ArtificialBeeColony:
    def __init__(self, dim, pop_size, iter_limit, obj_func, budget, limit, trial_limit):
        self.pop_size = pop_size
        self.obj_func = obj_func
        self.dim = dim
        self.limit = limit
        self.iter_limit = iter_limit
        self.tlim = trial_limit
        self.ebees = self.create_ebee()
        self.obees = self.create_obee()
        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, bees):
        for bee in bees:
            best_p = bee.position
            best_f = bee.fitness
            if best_f < self.gbest:
                self.gbest_p = deepcopy(best_p)
                self.gbest = deepcopy(best_f)

    def create_ebee(self):
        ebees = [EmployedBee(self.dim, self.limit, self.obj_func, self.tlim) for _ in range(self.pop_size // 2)]
        return ebees

    def create_obee(self):
        obees = [OnlookerBee(self.dim, self.limit, self.obj_func, self.tlim) for _ in range(self.pop_size // 2)]
        return obees

    def employee_phase(self):
        for bee in self.ebees:
            self.evaluations += bee.explore(self.ebees)

    def onlooker_phase(self):
        for bee in self.obees:
            self.evaluations += bee.exploit(self.ebees)

    def scout_phase(self):
        colony = self.ebees + self.obees
        for bee in colony:
            if bee.trial >= self.tlim:
                bee.position = np.random.uniform(self.limit[0], self.limit[1], self.dim)
                bee.fitness = self.obj_func(bee.position, self.dim)
                self.evaluations += 1
                bee.trial = 0
                bee.prob = 0

    def compute_probability(self):
        fitness_sum = 0
        for bee in self.ebees:
            fitness = bee.f.copy()
            fitness_sum += fitness

        for bee in self.ebees:
            bee.c_probability(fitness_sum)

    def evaluate(self):
        for iteration in range(self.iter_limit):
            if self.evaluations < self.budget:
                self.update_optimum(self.ebees)
                self.update_optimum(self.obees)
                self.gbest_history[iteration] = self.gbest
                self.gbest_history_p[iteration] = self.gbest_p
                self.employee_phase()
                self.update_optimum(self.ebees)
                self.compute_probability()
                self.onlooker_phase()
                self.update_optimum(self.obees)
                self.scout_phase()
                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}')