import numpy as np
import benchmark_functions
from copy import deepcopy

"""
    In constructor:
    Each particle has a set of attributes defined in constructor, the fitness_function is from benchmark_functions
    In velocity update:
    create two random indexes and use them to update the velocity according to equation, do this for every dimension
    defined, e.i - if dimension is 2 then velocity is defined with vector of 2D velocity [v1, v2], position as well
    Check the limits in check_Vlim -> if value exceeds limits, set them to upper or lower limit
    In position update:
    similiarly to velocity, update according to equation, then check limits, do this for every dimension
    Fitness_update:
    check if current fitness of particle is better than its personal best, if yes then upate it with its value
    comment: particle starts with its fitness and position equal to its personal best, however, personal best might or
    might not be updated, depends on the current fitness of particle -> which changes togehter with position in every
    iteration
"""


class Particle:
    def __init__(self, fitness_function, dim, pos_lim, vel_lim):
        self.vel_lim = vel_lim
        self.pos_lim = pos_lim
        self.dim = dim
        self.position = np.random.uniform(self.pos_lim[0], self.pos_lim[1], self.dim)
        self.velocity = np.random.uniform(self.vel_lim[0], self.vel_lim[1], self.dim)
        self.fitness_function = fitness_function
        self.fitness = self.fitness_function(self.position, self.dim)
        self.pbest_fitness = np.copy(self.fitness)
        self.pbest_pos = np.copy(self.position)

    def velocity_update(self, c1, c2, w, gbest):
        r1 = np.random.rand()
        r2 = np.random.rand()
        for d in range(self.dim):
            self.velocity[d] = (w * self.velocity[d] + c1 * r1 * (self.pbest_pos[d] - self.position[d])
                                + c2 * r2 * (gbest[d] - self.position[d]))
        self.check_Vlim()

    def position_update(self):
        for d in range(self.dim):
            self.position[d] = self.position[d] + self.velocity[d]
        self.check_Plim()

    def check_Vlim(self):
        for d in range(self.dim):
            if self.velocity[d] < self.vel_lim[0]:
                self.velocity[d] = self.vel_lim[0]
            elif self.velocity[d] > self.vel_lim[1]:
                self.velocity[d] = self.vel_lim[1]

    def check_Plim(self):
        for d in range(self.dim):
            if self.position[d] < self.pos_lim[0]:
                self.position[d] = self.pos_lim[0]
            elif self.position[d] > self.pos_lim[1]:
                self.position[d] = self.pos_lim[1]

    def fitness_update(self):
        self.fitness = self.fitness_function(self.position, self.dim)
        if self.fitness < self.pbest_fitness:
            self.pbest_fitness = self.fitness.copy()
            self.pbest_pos = self.position.copy()


"""
    In constructor
    Create particles in init_particles, each one with defined fitness function, dimension of velocity and position
    and position and velocity limits
    Set global fitness to zero, its position as well
    Set history of global optima to zero
    In evaluate:
    iterate for every particle and update its velocity and position and fitness then check if particle fitness
    is better than global best, if yes, rewrite
    after every generation of particles ( in iteration), save the global best of that iteration
    Init Particles:
    create a particle with parsed parameters, create as much particles as pop_size
    Get_g_best:
    for every particle in swarm, check its position and fitness, then copy them into two vector respectively
    after successful copy of position and fitness of each particle, check fitness vector for minimum fitness
    take index of that minimum fitness and save this fitness and position as global best using that index to
    take those best values out of two defined vectors
    Check_g_best:
    compare particle personal best ( the particle which is currently checked in iteration) and compare it with
    the global best, if its better than global, update the global with the value of personal
"""


class ParticleSwarmOptimization:
    def __init__(self, dim, pop_size, iter_limit, fitness_function, budget, pos_lim, vel_lim, c1, c2, w):
        self.dim = dim
        self.pop_size = pop_size
        self.c1 = c1
        self.c2 = c2
        self.w = w
        self.iter_limit = iter_limit
        self.particle_swarm = self.init_particles(fitness_function, self.dim, pos_lim, vel_lim)
        self.g_best = np.inf
        self.g_best_p = np.zeros([self.dim, 1])
        self.gbest_history = np.zeros([self.iter_limit, 1])
        self.gbest_history_p = np.zeros([self.iter_limit, self.dim])
        # every individual gets evaluated when created, thus we need to consider this evaluations
        self.evaluations = 0
        self.budget = budget
        self.iter_eval = [[-1, self.evaluations]]

    def init_particles(self, fitness_function, dim, pos_lim, vel_lim):
        swarm = [Particle(fitness_function, dim, pos_lim, vel_lim) for particle in range(self.pop_size)]
        return swarm

    def update_optimum(self):
        for particle in self.particle_swarm:
            best_p = particle.pbest_pos
            best_f = particle.pbest_fitness
            if best_f < self.g_best:
                self.g_best_p = deepcopy(best_p)
                self.g_best = deepcopy(best_f)

    def evaluate(self):
        self.update_optimum()
        for iteration in range(self.iter_limit):
            if self.evaluations < self.budget:
                for particle in self.particle_swarm:
                    if self.evaluations > self.budget:
                        break
                    else:
                        particle.velocity_update(self.c1, self.c2, self.w, self.g_best_p)
                        particle.position_update()
                        particle.fitness_update()
                        self.evaluations += 1
                self.gbest_history[iteration] = self.g_best
                self.gbest_history_p[iteration] = self.g_best_p
                self.update_optimum()
                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.g_best} at position {self.g_best_p}')
        print(f'History of global optima {self.gbest_history}')
        print(f'Evaluations {self.evaluations}')

