import benchmark_functions
import numpy as np
from copy import deepcopy

"""
    Moth is located at randomly generated position in specified limits, from passing this position as argument to 
    objective function fitness is calculated, apart from these variables, every moth carries variables D -> which is 
    distance between moth and flame, and t which is an coefficent calculated from a ( another variable, which is 
    lineary decreasing with iterations ) and a random number
    Moth updates its position according to equation defined by author, moving with respect to position of corresponding
    flame
    After moth moves, new position limits are checked and fitness is updated
"""


class Moth:
    def __init__(self, limits, obj_function, dim, b):
        self.lim = limits
        self.obj_function = obj_function
        self.dim = dim
        self.b = b
        self.position = np.random.uniform(self.lim[0], self.lim[1], self.dim)
        self.fitness = self.obj_function(self.position, self.dim)
        self.D = 0
        self.t = 0

    def calculate_distance(self, flame):
        self.D = np.abs(flame.position - self.position)

    def calculate_t(self, a):
        self.t = (a - 1) * np.random.uniform(0, 1, self.dim) + 1

    def move(self, flame, a):
        self.calculate_distance(flame)
        self.calculate_t(a)
        self.position = self.D * np.exp(self.b * self.t) * np.cos(self.t * 2 * np.pi) + flame.position
        # for d in range(self.dim):
        #     self.position[d] = self.D[d] * np.exp(self.b * self.t[d]) * np.cos(self.t[d] * 2 * np.pi) + flame.position[d]
        self.check_lim()
        self.update_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 update_fitness(self):
        self.fitness = self.obj_function(self.position, self.dim)


"""
    In MFO, flames and moths are the same object, but flames represent the best moth positions
    First initial population of moths is created, then these are copied into flames ( which will be sorted afterwards)
    Evaluate: at the beginning, flames are sorted, then, for each iteration, a and flame_no is counted. These two 
    parameters are calculated by equations defined by author, a is vital for movement of moth, flame_no is number of
    flame, which is going to be implemented in moth movement with respect to current iteration, this operation was 
    also defined by author
    In first iteration, we use sorted flames for movement of moths, every moth is moving with respect to its corresponding
    flame ( corresponding with respect to index -> i ) OR movement is computed with flame on index flames_no, if certain
    condition is met ( that i > flames_no)
    Moths change their position and update their fitness, however. optimum is updated from flames
    In every other iteration, flames are updated before updating position of moths, their value is updated from a list,
    which contains joint population of moths and flames and is sorted. Then, new population of flames, the updated 
    population, is taken from this sorted list of joint flames and moths ( we take pop_size number of objects )
    This updated flames are then used in computation of movement of moths, and after their update, we select global
    optima from this list and update history as well
"""


class MothFlameOptimization:
    def __init__(self, dimension, pop_size, iter_limit, obj_function, budget, limits, b):
        self.b = b
        self.dimension = dimension
        self.limits = limits
        self.obj_function = obj_function
        self.pop_size = pop_size
        self.iter_limit = iter_limit
        self.moth_swarm = self.create_moths()
        self.flames = self.create_flames()
        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.a = None
        self.b = b
        self.flame_no = 0
        self.budget = budget
        self.evaluations = 0
        self.iter_eval = [[-1, self.evaluations]]

    def create_moths(self):
        moths = [Moth(self.limits, self.obj_function, self.dimension, self.b) for _ in range(self.pop_size)]
        return moths

    def create_flames(self):
        flames = [deepcopy(moth) for moth in self.moth_swarm]
        return flames

    def compute_flame_no(self, iteration):
        self.flame_no = round(self.pop_size - (iteration + 1) * ((self.pop_size - 1) / self.iter_limit))
        if self.flame_no == self.pop_size:
            self.flame_no -= 1

    def update_optimum(self, moth_flame):
        #self.gbest_p = moth_flame[0].position
        #self.gbest = moth_flame[0].fitness
        for mfo in moth_flame:
            best_p = mfo.position
            best_f = mfo.fitness
            if best_f < self.gbest:
                self.gbest_p = deepcopy(best_p)
                self.gbest = deepcopy(best_f)

    def reduce_a(self, iteration):
        self.a = -1 + (iteration + 1) * ((-1) / self.iter_limit)

    def sort(self, moth_flame):
        moth_flame.sort(key=lambda mof: mof.fitness)

    def update_flames(self):
        positions = self.moth_swarm + self.flames
        positions.sort(key=lambda pos: pos.fitness)
        for idx in range(self.pop_size):
            self.flames[idx] = deepcopy(positions[idx])

    def update_optimum_history(self, iteration):
        self.gbest_history[iteration] = self.gbest
        self.gbest_history_p[iteration] = self.gbest_p

    def evaluate(self):
        self.sort(self.flames)
        self.update_optimum(self.flames)
        for iteration in range(self.iter_limit):
            if self.evaluations < self.budget:
                self.reduce_a(iteration)
                self.compute_flame_no(iteration)
                i = 0
                if iteration == 0:
                    pass
                else:
                    self.update_flames()
                for moth in self.moth_swarm:
                    if self.evaluations > self.budget:
                        break
                    if iteration <= self.flame_no:
                        moth.move(self.flames[i], self.a)
                        self.evaluations += 1
                    else:
                        moth.move(self.flames[self.flame_no], self.a) #self.gbest ?
                        self.evaluations += 1
                    i += 1
                self.update_optimum_history(iteration)
                self.update_optimum(self.flames)
                self.iter_eval.append([iteration, self.evaluations])
                # self.print_partial(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}')

    def print_partial(self, iteration):
        if (iteration % 100 == 0):
            print(f'Fitness {self.gbest} at iteration {iteration}')


