import random
import numpy as np
from bejoor.core.bejoor_algorithm import BejoorAlgorithm
from bejoor.core.individual import Individual

class FireworksAlgorithm(BejoorAlgorithm):
    def __init__(self, *args, explosion_strength=50, sparks_per_firework=40, **kwargs):
        """
        Fireworks Algorithm.
        :param explosion_strength: Strength of the explosion, determining the distance sparks travel.
        :param sparks_per_firework: Number of sparks generated by each firework.
        """
        super().__init__(*args, **kwargs)
        self.explosion_strength = explosion_strength
        self.sparks_per_firework = sparks_per_firework
        self.optimizer_name = "Fireworks Algorithm"

    def explode(self, firework):
        """Explode a firework to generate sparks around it."""
        sparks = []
        for _ in range(self.sparks_per_firework):
            new_solution = []
            for i in range(self.solution_vector_size):
                if random.random() < 0.5:
                    new_value = firework.values[i] + random.uniform(-self.explosion_strength, self.explosion_strength)
                else:
                    new_value = firework.values[i] - random.uniform(-self.explosion_strength, self.explosion_strength)

                # Keep the values within the solution bounds
                new_value = max(self.solution_vector[i]['lower_bound'], min(self.solution_vector[i]['upper_bound'], new_value))
                new_solution.append(new_value)

            # Create new spark individual
            spark = Individual(firework.id, new_solution)
            sparks.append(spark)

        return sparks

    def update(self):
        """Update the population by exploding fireworks and generating sparks."""
        new_population = []

        # Explode each firework to generate sparks
        for firework in self.population:
            sparks = self.explode(firework)
            new_population.extend(sparks)

        # Evaluate the objective values for the new population (sparks)
        self.population = new_population
        self.evaluate_all_objectives()  # Ensure all individuals have their objective values set

        # Sort and keep only the best individuals
        self.population = sorted(self.population, key=lambda ind: ind.get_objective_value())[:self.population_size]

        # Update the best solution
        self.best_solution = self.population[0].values
        self.best_objective_value = self.population[0].objective_value

        # Update global best if necessary
        if ((self.global_best_objective_value < self.best_objective_value and self.optimization_side == 'max') or
                (self.global_best_objective_value > self.best_objective_value and self.optimization_side == 'min')):
            self.global_best_solution = self.best_solution
            self.global_best_objective_value = self.best_objective_value
