from scipy.spatial.distance import cdist
from search.MOPSO import *
from random import random


class BBO(MOEA):
    def __init__(self, proteins, config, energy, coder, logging, current_gen=1):
        super(BBO, self).__init__(config, energy, coder, logging, current_gen)
        self.proteins = proteins
        self.obj_num = self.proteins[0].obj_num
        self.max_angles, self.min_angles = proteins[0].get_angles_field()
        self.init_param()

    def init_param(self):
        self.KeepRate = self.config['KeepRate']  # The number of neighbors for estimating density
        self.alpha = self.config['alpha']  # The number of candidates for convergence-based selection
        self.pMutation = self.config['pMutation']
        self.archive_thresh = self.config['archive_thresh']
        self.nVar = self.config['nVar']

        self.nKeep = np.round(self.KeepRate * self.pop_size)
        self.nNew = self.pop_size - self.nKeep

        self.mu = np.linspace(1, 0, self.pop_size)  # Emmigration Rates
        self.lamb = 1 - self.mu  # Immigration Rates
        self.sigma = 0.02 * (self.max_angles - self.min_angles)

    def RouletteWheelSelection(self, EP):
        r = random()
        C = np.cumsum(EP)
        j = np.argwhere(r <= C)[0]
        return j

    def selection_protein(self):
        return

    def run(self):
        time1 = time.time()
        fitness = np.array([x.obj.copy() for x in self.proteins])
        population = np.array([x.angle_view().copy() for x in self.proteins])
        sort_index = np.argsort(fitness)
        population = population[sort_index]
        fitness = fitness[sort_index]
        self.proteins = self.proteins[sort_index]

        best_fitness_list = np.zeros((self.max_gen, 1))
        self.offspring = [x.copy() for x in self.proteins]

        for x in range(self.current_gen, self.max_gen):
            start_time = time.time()
            new_pop = population.copy()

            # evolution
            for i in range(len(new_pop)):
                for k in range(self.nVar):
                    if random() < self.lamb[i]:
                        EP = self.mu.copy()
                        EP[i] = 0
                        EP = EP / np.sum(EP)
                        j = self.RouletteWheelSelection(EP)

                        new_pop[i] = population[i] + self.alpha * (population[j, k] - population[i, k])

                        if random() <= self.pMutation:
                            new_pop[i, k] = new_pop[i, k] + self.sigma * random()

            for protein, angle in zip(self.offspring, new_pop):
                protein.update_angle_from_view(angle)

            new_fitness = np.array([x.obj.copy() for x in self.offspring])
            new_sort_index = np.argsort(new_fitness)
            new_pop = new_pop[new_sort_index]
            new_fitness = new_fitness[new_sort_index]

            self.offspring = self.offspring[new_sort_index]

            combine_fitness = (fitness[self.nKeep] + new_fitness[self.nNew]).copy()
            final_index = np.argsort(combine_fitness)
            combine_fitness = combine_fitness[final_index]
            self.proteins = (self.proteins[self.nKeep] + self.offspring[self.nNew])[final_index].copy()
            best_fitness_list[x] = combine_fitness[0]
            print('The current best fitenss:', combine_fitness[0])
            self.current_gen = x
            self.logging.write(self.proteins, self.coder, self.config['save_all'], self.current_gen)
        f = open('BBO_fitness.txt', 'w')
        for x in best_fitness_list:
            f.write(x + '\n')
        f.close()
