"""
Class for running optimization experiments using Genetic Algorithms (GA), including grid search functionality.

Example usage:
    experiment_name = 'example_experiment'
    problem = TSPGenerator.generate(seed=SEED, number_of_cities=22)

    ga = GARunner(problem=problem,
                  experiment_name=experiment_name,
                  output_directory=OUTPUT_DIRECTORY,
                  seed=SEED,
                  iteration_list=2 ** np.arange(12),
                  max_attempts=1000,
                  population_sizes=[150, 200, 300],
                  mutation_rates=[0.4, 0.5, 0.6])

    df_run_stats, df_run_curves = ga.run()
"""

# Authors: Andrew Rollings (modified by Kyle Nakamura)
# License: BSD 3-clause

from typing import Any

import numpy as np
import pandas as pd

from mlrose_ky.algorithms import genetic_alg
from mlrose_ky.decorators import short_name
from mlrose_ky.runners._runner_base import _RunnerBase


@short_name("ga")
class GARunner(_RunnerBase):
    """
    A runner for performing optimization experiments using Genetic Algorithms (GA).

    This class extends _RunnerBase and provides functionality for running experiments
    with Genetic Algorithms, including grid search over hyperparameters such as
    population size and mutation rate.

    Attributes
    ----------
    population_sizes : list[int]
        List of population sizes to test in the grid search.
    mutation_rates : list[float]
        List of mutation rates to test in the grid search.
    hamming_factors : Optional[list[float]]
        Optional list of Hamming factors for the algorithm.
    hamming_factor_decays : Optional[list[float]]
        Optional list of Hamming factor decays.
    """

    def __init__(
        self,
        problem: Any,
        experiment_name: str,
        seed: int,
        iteration_list: np.ndarray | list[int],
        population_sizes: list[int],
        mutation_rates: list[float],
        hamming_factors: list[float] = None,
        hamming_factor_decays: list[float] = None,
        max_attempts: int = 500,
        generate_curves: bool = True,
        output_directory: str = None,
        **kwargs: Any,
    ):
        """
        Initialize the GARunner class with problem data and various experiment parameters.

        Parameters
        ----------
        problem : Any
            The optimization problem to be solved.
        experiment_name : str
            Name of the experiment.
        seed : int
            Random seed for reproducibility.
        iteration_list : np.ndarray | list of int
            List of iterations for the experiment.
        population_sizes : list of int
            List of population sizes to test in the grid search.
        mutation_rates : list of float
            List of mutation rates to test in the grid search.
        hamming_factors : list of float, optional
            Optional list of Hamming factors for the algorithm.
        hamming_factor_decays : list of float, optional
            Optional list of Hamming factor decays.
        max_attempts : int, optional
            Maximum number of attempts without improvement before stopping.
        generate_curves : bool, optional
            Whether to generate learning curves.
        output_directory : str, optional
            Directory to save experiment result, default=None.
        """
        super().__init__(
            problem=problem,
            experiment_name=experiment_name,
            seed=seed,
            iteration_list=iteration_list,
            max_attempts=max_attempts,
            generate_curves=generate_curves,
            output_directory=output_directory,
            **kwargs,
        )
        self.population_sizes: list[int] = population_sizes
        self.mutation_rates: list[float] = mutation_rates
        self.hamming_factors: list[float] = hamming_factors
        self.hamming_factor_decays: list[float] = hamming_factor_decays

    def run(self) -> tuple[pd.DataFrame, pd.DataFrame]:
        """
        Run the genetic algorithm experiment.

        This method performs grid search over the provided population sizes,
        mutation rates, and optional Hamming factors and decays, and returns
        the statistics and curves generated by the experiment.

        Returns
        -------
        tuple
            A tuple containing two DataFrames: run statistics and run curves
        """
        return super().run_experiment_(
            algorithm=genetic_alg,
            pop_size=("Population Size", self.population_sizes),
            mutation_prob=("Mutation Rate", self.mutation_rates),
            hamming_factor=("Hamming Factor", self.hamming_factors),
            hamming_decay_factor=("Hamming Factor Decay Rate", self.hamming_factor_decays),
        )
