import numpy as np
import benchmark_functions
from copy import deepcopy

"""
    Elephant class has several attributes, that is alpha and beta -> which are control parameters, then  objective
    function dimension and limits, and finally its position -> randomly generated and fitness, which is obtained by
    passing position to objective function
    Elephant can move by two rules, by which one depends if its matriarch-> best elephant in clan
    Move_to_matriarch -> this movement is utilized by every elephant which is not matriarch, the new position is 
    computed by using alpha parameter, random number r and difference between best elefant -> matriarch and position
    of elephant currently updating its position
    Move_matriach -> this movement is utilized by matriarch -> the best elephant in clan, so its new position is 
    determined by pamrameter beta and clan center -> this is computed by arithmetic mean of all positions of elephants
    (dimension by dimension respectively) = first all positions are taken from elephants, then mean is computed
    then new position is updated with respect to each dimension
    Check_dimension serves for checking the limits of new position
"""


class Elephant:
    def __init__(self, alpha, beta, obj_func, dim, limits):
        self.alpha = alpha
        self.beta = beta
        self.obj_func = obj_func
        self.dim = dim
        self.lim = limits
        self.position = np.random.uniform(self.lim[0], self.lim[1], self.dim)
        self.fitness = self.obj_func(self.position, self.dim)

    def update_fitness(self):
        self.fitness = self.obj_func(self.position, self.dim)

    def move(self, is_matriarch, clan):
        if is_matriarch:
            self.move_matriarch(clan)
        else:
            self.move_to_matriarch(clan)

    def move_to_matriarch(self, clan):
        matriarch = clan[0]
        for d in range(self.dim):
            r = np.random.uniform()
            self.position[d] = self.position[d] + self.alpha * r * (matriarch.position[d] - self.position[d])
        self.check_lim()
        self.update_fitness()

    def move_matriarch(self, clan):
        positions = np.array([elephant.position for elephant in clan])
        clan_center = np.mean(positions, axis=0)
        for d in range(self.dim):
            self.position[d] = self.beta * clan_center[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]


"""
    ElephantHerdOptimization utilizes several parameters, two of them alpha and beta are control parameters
    In constructor, elephant clans are created => each clan consists of elephants of defined number clan_size
    this number is determined from population size and number of clans
    Sort_elephants is sorting elephants in clan according to their fitness, the first elephant in clan is therefore the 
    best and is labeled matriarch
    Evaluate function is running for iter_limit times, while calling two important methods specific to EHO those are:
    Clan_update -> this method is implementing movement of elephant, so for each clan a elephant is selected, the first
    one is matriarch, therefore after first iteration, the condition is set to false -> movement is utilized by 
    elephants respectively according to their status (if they are matriarch or not)
    After this update, we sort elephants and update optimum with hope of finding better solution, then another specific
    method is called
    Separation -> this method gets rid of worst elephant in clan, replacing him by another solution (elephant) with 
    new position and fitness
    Again after this method we sort elephants and update optimum, to check if new best solution appeared ( from new 
    elephants )
"""


class ElephantHerdOptimization:
    def __init__(self, dimension, pop_size, iter_limit, obj_function, budget, limits, alpha, beta, no_clans):
        self.alpha = alpha
        self.beta = beta
        self.dimension = dimension
        self.limits = limits
        self.objective_function = obj_function
        self.no_clans = no_clans
        self.pop_size = pop_size
        self.clan_size = self.pop_size // self.no_clans
        self.iter_limit = iter_limit
        self.elephant_clans = self.create_clans()
        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.budget = budget
        self.evaluations = 0
        self.iter_eval = [[-1, self.evaluations]]

    def create_clans(self):
        clans = []
        for clan in range(self.no_clans):
            elephants = [Elephant(self.alpha, self.beta,self.objective_function, self.dimension, self.limits)
                         for _ in range(self.clan_size)]
            clans.append(elephants)
        return clans

    def sort_elephants(self):
        for clan in self.elephant_clans:
            clan.sort(key=lambda elephant: elephant.fitness)

    def update_optimum(self):
        for clan in self.elephant_clans:
            best_p = clan[0].position
            best_f = clan[0].fitness
            if best_f < self.gbest:
                self.gbest_p = deepcopy(best_p)
                self.gbest = deepcopy(best_f)

    def clan_update(self):
        for clan in self.elephant_clans:
            is_matriarch = True
            for elephant in clan:
                elephant.move(is_matriarch, clan)
                is_matriarch = False
                self.evaluations += 1


    def separate(self):
        for clan in self.elephant_clans:
            clan[-1] = Elephant(self.alpha, self.beta, self.objective_function, self.dimension, self.limits)
            self.evaluations += 1

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

    def evaluate(self):
        self.update_optimum()
        for iteration in range(self.iter_limit):
            if self.evaluations < self.budget:
                self.update_optimum_history(iteration)
                self.sort_elephants()
                self.clan_update()
                self.sort_elephants()
                self.update_optimum()
                self.separate()
                self.update_optimum()
                self.sort_elephants()
                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.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}')

