# evolution/algorithm.py
import copy
import enum
import logging
import random
import time
from typing import List, Set, Tuple

import deprecation

from fandango.constraints.base import Constraint
from fandango.constraints.fitness import FailingTree, Comparison, ComparisonSide
from fandango.language.grammar import DerivationTree
from fandango.language.grammar import Grammar
from fandango.logger import LOGGER, visualize_evaluation, clear_visualization


class LoggerLevel(enum.Enum):
    NOTSET = logging.NOTSET
    DEBUG = logging.DEBUG
    INFO = logging.INFO
    WARNING = logging.WARNING
    ERROR = logging.ERROR
    CRITICAL = logging.CRITICAL


class Fandango:
    def __init__(
        self,
        grammar: Grammar,
        constraints: List[Constraint],
        population_size: int = 100,
        desired_solutions: int = 0,
        initial_population: List[DerivationTree | str] = None,
        max_generations: int = 500,
        elitism_rate: float = 0.1,
        crossover_rate: float = 0.8,
        tournament_size: float = 0.1,
        mutation_rate: float = 0.2,
        destruction_rate: float = 0.0,
        logger_level: LoggerLevel = None,
        warnings_are_errors: bool = False,
        best_effort: bool = False,
        random_seed: int = None,
        start_symbol="<start>",
    ):
        """
        Initialize the FANDANGO genetic algorithm. The algorithm will evolve a population of individuals
        that are generated by the given grammar, and will evaluate them against the given constraints.

        :param grammar: The grammar used to generate individuals.
        :param constraints: The constraints used to evaluate individuals.
        :param population_size: The size of the population.
        :param desired_solutions: The number of perfect solutions to find before stopping the algorithm.
        :param initial_population: A list of individuals to use as the initial population.
        :param max_generations: The maximum number of generations to run the algorithm.
        :param elitism_rate: The rate of individuals that will be preserved in the next generation.
        :param crossover_rate: The rate of individuals that will undergo crossover.
        :param mutation_rate: The rate of individuals that will undergo mutation.
        :param tournament_size: The size of the tournament selection.
        :param destruction_rate: The rate of individuals that will be destroyed.
        :param logger_level: If set, the level of logging to use. One of DEBUG, INFO, WARNING, ERROR, CRITICAL.
        :param start_symbol: The start symbol to use with the grammar.
        :param warnings_are_errors: If set, turns warnings into errors
        :param best_effort: If set, returns also solutions not satisfying all constraints
        :param random_seed: The random seed to use for reproducibility.
        """
        if tournament_size > 1:
            raise ValueError(
                f"Parameter tournament_size must be in range ]0, 1], but is {tournament_size}."
            )
        if random_seed is not None:
            random.seed(random_seed)

        if logger_level is not None:
            LOGGER.setLevel(logger_level.value)

        LOGGER.info(f"---------- Initializing FANDANGO algorithm ---------- ")
        self.grammar = grammar
        self.constraints = constraints
        self.population_size = population_size
        self.mutation_rate = mutation_rate
        self.crossover_rate = crossover_rate
        self.tournament_size = max(2, int(population_size * tournament_size))
        self.max_generations = max_generations
        self.elitism_rate = elitism_rate
        self.destruction_rate = destruction_rate
        self.start_symbol = start_symbol

        self.fitness_cache = {}

        self.fixes_made = 0
        self.checks_made = 0
        self.crossovers_made = 0
        self.mutations_made = 0

        self.time_taken = None

        self.warnings_are_errors = warnings_are_errors
        self.best_effort = best_effort

        # Initialize population
        self.solution = list()
        self.desired_solutions = desired_solutions

        if initial_population is not None:
            LOGGER.info(f"Saving the provided initial population...")
            self.population = []
            for individual in initial_population:
                if isinstance(individual, str):
                    tree = self.grammar.parse(individual)
                    if not tree:
                        raise ValueError(
                            f"Failed to parse initial individual: {individual}"
                        )
                    self.population.append(tree)
                elif isinstance(individual, DerivationTree):
                    self.population.append(individual)
                else:
                    raise TypeError(
                        f"Initial individuals must be DerivationTree or String"
                    )
            for i in range(self.population_size - len(self.population)):
                self.population.append(self.mutate(self.population[i]))
        else:
            LOGGER.info(
                f"Generating initial population (size: {self.population_size})..."
            )

            st_time = time.time()
            self.population = self.generate_random_initial_population()
            LOGGER.info(
                f"Initial population generated in {time.time() - st_time:.2f} seconds"
            )

        # Evaluate population
        self.evaluation = self.evaluate_population()
        self.fitness = (
            sum(fitness for _, fitness, _ in self.evaluation) / self.population_size
        )

    def evolve(self) -> List[DerivationTree]:
        """
        Run the genetic algorithm to evolve the population over multiple generations.

        :return: The best solution found by the algorithm.
        """
        LOGGER.info(f"---------- Starting evolution ----------")

        start_time = time.time()

        for generation in range(1, self.max_generations + 1):
            if 0 < self.desired_solutions <= len(self.solution):
                self.fitness = 1.0
                self.solution = self.solution[: self.desired_solutions]
                break
            if len(self.solution) >= self.population_size:
                self.fitness = 1.0
                self.solution = self.solution[: self.population_size]
                break
            if self.fitness >= 0.99:
                self.fitness = 1.0
                self.solution = self.population[: self.population_size]
                break

            LOGGER.info(
                f"Generation {generation} - Fitness: {self.fitness:.2f} - "
                f"#solutions found: {len(self.solution)}"
            )

            # Select elites
            new_population = self.select_elites()

            # Crossover
            while len(new_population) < self.population_size:
                if random.random() < self.crossover_rate:
                    parent1, parent2 = self.tournament_selection()
                    child1, child2 = self.crossover(parent1, parent2)
                    new_population.append(child1)
                    new_population.append(child2)
                    self.crossovers_made += 1

            new_population = new_population[: self.population_size]

            # Mutation
            for individual in new_population:
                if random.random() < self.mutation_rate:
                    new_population.remove(individual)
                    new_population.append(self.mutate(individual))
                    self.mutations_made += 1

            # Destruction
            if self.destruction_rate > 0:
                LOGGER.debug(
                    f"Destroying {self.destruction_rate * 100:.2f}% of the population"
                )
                random.shuffle(new_population)
                new_population = new_population[
                    : int(self.population_size * (1 - self.destruction_rate))
                ]

            # Add new individuals
            while len(new_population) < self.population_size:
                new_population.append(self.grammar.fuzz(start=self.start_symbol))

            # Fix individuals
            fixed_population = list()
            for individual in new_population:
                fixed_population.append(self.fix_individual(individual))

            # Evaluate population
            self.population = fixed_population[: self.population_size]
            self.evaluation = self.evaluate_population()
            self.fitness = (
                sum(fitness for _, fitness, _ in self.evaluation) / self.population_size
            )

            # Report results
            visualize_evaluation(generation, self.max_generations, self.evaluation)

        clear_visualization()
        self.time_taken = time.time() - start_time

        LOGGER.info(f"---------- Evolution finished ----------")
        LOGGER.info(f"Perfect solutions found: ({len(self.solution)})")
        LOGGER.info(f"Fitness of final population: {self.fitness:.2f}")
        LOGGER.info(f"Time taken: {self.time_taken:.2f} seconds")

        LOGGER.debug(f"---------- FANDANGO statistics ----------")
        LOGGER.debug(f"Fixes made: {self.fixes_made}")
        LOGGER.debug(f"Fitness checks: {self.checks_made}")
        LOGGER.debug(f"Crossovers made: {self.crossovers_made}")
        LOGGER.debug(f"Mutations made: {self.mutations_made}")

        if self.fitness < 1.0:
            LOGGER.error(f"Population did not converge to a perfect population")
            if self.warnings_are_errors:
                raise RuntimeError("Failed to find a perfect solution")
            if self.best_effort:
                return self.population

        if self.desired_solutions > 0 and len(self.solution) < self.desired_solutions:
            LOGGER.error(
                f"Only found {len(self.solution)} perfect solutions, instead of the required {self.desired_solutions}"
            )
            if self.warnings_are_errors:
                raise RuntimeError(
                    "Failed to find the required number of perfect solutions"
                )
            if self.best_effort:
                return self.population[: self.desired_solutions]

        return self.solution

    def generate_random_initial_population(self) -> List[DerivationTree]:
        """
        Generate the initial population of individuals.

        :return: A set of individuals.
        """

        population = [
            self.grammar.fuzz(self.start_symbol) for _ in range(self.population_size)
        ]

        # Fix individuals
        fixed_population = list()
        for individual in population:
            fixed_population.append(self.fix_individual(individual))

        return fixed_population

    def fix_individual(self, individual: DerivationTree) -> DerivationTree:
        """
        Fix an individual by replacing failing subtrees if ComparisonConstraint.EQUAL are involved.
        """

        evaluation = self.evaluate_individual(individual)
        failing_trees = evaluation[1]
        for failing_tree in failing_trees:
            for operator, value, side in failing_tree.suggestions:
                if operator == Comparison.EQUAL and side == ComparisonSide.LEFT:
                    suggested_tree = self.grammar.parse(
                        str(value), failing_tree.tree.symbol
                    )
                    if suggested_tree is None:
                        continue
                    individual = individual.replace(failing_tree.tree, suggested_tree)
                    self.fixes_made += 1
        return individual

    def evaluate_individual(
        self, individual: DerivationTree
    ) -> Tuple[float, List[FailingTree]]:
        """
        Evaluate the fitness of an individual.

        :param individual: The individual to evaluate.
        :return: The fitness of the individual and the list of failing trees.
        """
        fitness = 0.0
        failing_trees = []

        if str(individual) in self.fitness_cache:
            if self.fitness_cache[str(individual)][0] >= 0.99:
                self.solution.append(individual)
            return self.fitness_cache[str(individual)]

        for constraint in self.constraints:
            result = constraint.fitness(individual)
            if result.success:
                fitness += result.fitness()
            else:
                failing_trees.extend(result.failing_trees)
                fitness += result.fitness()
            self.checks_made += 1

        # Normalize fitness
        try:
            fitness /= len(self.constraints)
        except ZeroDivisionError:
            fitness = 1.0

        if fitness >= 0.99:
            self.solution.append(individual)

        self.fitness_cache[str(individual)] = [fitness, failing_trees]
        return fitness, failing_trees

    def evaluate_population(
        self,
    ) -> List[Tuple[DerivationTree, float, List[FailingTree]]]:
        """
        Evaluate the fitness of each individual in the population.

        :return: A list of tuples, each containing an individual, its fitness, and the list of failing trees.
        """
        evaluation = []
        for individual in self.population:
            fitness, failing_trees = self.evaluate_individual(individual)
            evaluation.append((individual, fitness, failing_trees))
        return evaluation

    def select_elites(self) -> List[DerivationTree]:
        """
        Select the 'elitism'% elite individuals from the population based on the fitness.

        :return: A list of elite individuals.
        """
        return [
            x[0]
            for x in sorted(self.evaluation, key=lambda x: x[1], reverse=True)[
                : int(self.elitism_rate * self.population_size)
            ]
        ]

    def tournament_selection(self) -> Tuple[DerivationTree, DerivationTree]:
        """
        Perform tournament selection to choose two parents for crossover.
        """
        tournament = random.sample(self.evaluation, k=self.tournament_size)
        tournament.sort(key=lambda x: x[1], reverse=True)
        parent1 = tournament[0][0]
        parent2 = tournament[1][0]
        return parent1, parent2

    # noinspection PyMethodMayBeStatic
    def crossover(
        self, parent1: DerivationTree, parent2: DerivationTree
    ) -> Tuple[DerivationTree, DerivationTree]:
        """
        Perform crossover between two parents to generate two children by swapping subtrees rooted at a common
        non-terminal symbol.
        """
        # Get all non-terminal symbols in parent1 and parent2
        symbols1 = parent1.get_non_terminal_symbols()
        symbols2 = parent2.get_non_terminal_symbols()

        # Find common non-terminal symbols
        common_symbols = symbols1.intersection(symbols2)

        if not common_symbols:
            return parent1, parent2

        # Randomly select a common non-terminal symbol
        symbol = random.choice(list(common_symbols))

        # Find all nodes with that symbol in parent1 and parent2
        nodes1 = parent1.find_all_nodes(symbol)
        nodes2 = parent2.find_all_nodes(symbol)

        # Randomly select one node from each parent
        node1 = random.choice(nodes1)
        node2 = random.choice(nodes2)

        # Swap subtrees
        child1 = parent1.replace(node1, copy.deepcopy(node2))
        child2 = parent2.replace(node2, copy.deepcopy(node1))

        return child1, child2

    def mutate(self, individual: DerivationTree) -> DerivationTree:
        """
        Perform mutation on an individual to generate a new individual based on the failing trees.
        """
        failing_trees = self.evaluate_individual(individual)[1]
        selection = []

        for failing_tree in failing_trees:
            selection.append(failing_tree.tree)

        if len(selection) == 0:
            return individual
        else:
            node_to_mutate = random.choice(selection)
            if node_to_mutate.symbol.is_non_terminal:
                new_subtree = self.grammar.fuzz(node_to_mutate.symbol)
                individual = individual.replace(node_to_mutate, new_subtree)
                self.mutations_made += 1
        return individual


# Backwards compatibility
class FANDANGO(Fandango):
    @deprecation.deprecated(details="Use `Fandango` instead")
    def __init__(*args, **kwargs):
        super().__init__(*args, **kwargs)
