import random

from tona_ai.neat.environment import Environment
from tona_ai.neat.genome import Genome
from tona_ai.neat.individual import Individual
from tona_ai.neat.population import Population


class NEAT:
    """
    Represents the NEAT (NeuroEvolution of Augmenting Topologies) algorithm for evolving a population of neural networks.

    Attributes:
        population (Population): The population of individuals (neural networks) being evolved.
        environment (Environment): The environment in which the individuals are evaluated.
        mutation_rate (float): The rate at which mutations occur in individuals. Default is 0.1.
        mutation_range (tuple[float, float]): The range within which mutation values are applied. Default is (-0.1, 0.1).
        fitness_history (list[float]): A history of the best fitness scores for each epoch.

    Methods:
        run(epochs: int) -> Population:
            Runs the evolution process for the specified number of epochs.
        evolve():
            Evolves the population by selecting the best individuals and creating new individuals through reproduction and mutation.
        mutate(individual: Individual):
            Mutates the given individual by adjusting weights and biases randomly.
        mutate_population():
            Mutates the entire population by applying mutations to each individual.
    """

    def __init__(
        self,
        population: Population,
        environment: Environment,
        mutation_rate: float = 0.1,
        mutation_range: tuple[float, float] = (-0.1, 0.1),
    ):
        """
        Initializes the NEAT algorithm with a population, environment, mutation rate, and mutation range.

        Args:
            population (Population): The population of individuals (neural networks) to evolve.
            environment (Environment): The environment used to evaluate the fitness of individuals.
            mutation_rate (float, optional): The probability that a mutation will occur in each individual. Default is 0.1.
            mutation_range (tuple[float, float], optional): The range within which mutation values (weights and biases) will be applied. Default is (-0.1, 0.1).
        """
        self.population = population
        self.environment = environment
        self.mutation_rate = mutation_rate
        self.mutation_range = mutation_range

        self.fitness_history = []

    def run(self, epochs: int) -> Population:
        """
        Runs the NEAT algorithm for the specified number of epochs, evaluating and evolving the population each epoch.

        Args:
            epochs (int): The number of epochs (iterations) to run the NEAT algorithm.

        Returns:
            Population: The evolved population after the specified number of epochs.
        """
        for i in range(epochs):
            # Evaluate the population's fitness in the given environment
            self.fitness_history.append(self.population.evaluate(self.environment))
            best_fitness = max(ind.fitness for ind in self.population.individuals)
            print(f"Epoch {i}, Best Fitness: {best_fitness}")
            self.evolve()
        return self.population

    def evolve(self):
        """
        Evolves the population by selecting the top-performing individuals and creating new individuals through
        reproduction (crossover) and mutation.

        This method creates a new generation by keeping the top 50% of individuals and replacing the bottom 50% with
        offspring generated by crossover and mutation.
        """
        sorted_individuals = sorted(
            self.population.individuals, key=lambda x: x.fitness
        )

        # Select the top 50% and bottom 50% of individuals
        top_50_percent = sorted_individuals[: len(self.population.individuals) // 2]
        bottom_50_percent = sorted_individuals[len(self.population.individuals) // 2 :]

        # Reproduce and mutate the bottom 50% using individuals from the top 50%
        for i in range(len(bottom_50_percent)):
            child_genome = Genome(
                top_50_percent[i].genome.neurons, top_50_percent[i].genome.synapses
            )
            child = Individual(genome=child_genome)
            self.mutate(child)
            bottom_50_percent[i] = child

    def mutate(self, individual: Individual):
        """
        Mutates the given individual by adjusting its synapse weights and neuron biases randomly.

        Args:
            individual (Individual): The individual to mutate.
        """
        # Mutate synapses weights
        if random.random() < self.mutation_rate:
            for synapse in individual.genome.synapses:
                if random.random() < self.mutation_rate:
                    synapse.weight += random.uniform(
                        self.mutation_range[0], self.mutation_range[1]
                    )

        # Mutate neuron biases
        if random.random() < self.mutation_rate:
            for neuron in individual.genome.neurons:
                if random.random() < self.mutation_rate:
                    neuron.bias += random.uniform(
                        self.mutation_range[0], self.mutation_range[1]
                    )

    def mutate_population(self):
        """
        Mutates the entire population by applying mutations to each individual.

        This method is typically used after evolving the population to apply small changes to all individuals,
        helping introduce diversity and potentially improve the population's performance.

        Returns:
            Population: The population after mutations have been applied.
        """
        for individual in self.population.individuals:
            self.mutate(individual)

        return self.population
