from typing import Tuple, Dict
import random

from gpforecaster.model.gpf import GPF
from gpforecaster.utils.logger import Logger


def optimize_hyperparameters(
    dataset_name: str,
    hierarchical_data: Dict,
    num_trials: int = 20,
    learning_rates: Tuple[float] = (1e-2, 1e-3),
    weight_decays: Tuple[float] = (1e-3, 1e-4, 1e-5),
    scheduler_types: Tuple[str] = ("step", "exponential", "cosine", "none"),
    gamma_rates: Tuple[int] = (0.1, 0.9, 0.95, 0.8),
) -> Tuple[float, float, str, float]:
    """
    Performs a random search to optimize the hyperparameters of the model.

    Parameters
        dataset: Dataset to run hyperparameter tuning
        hierarchical_data: Dict containing data and metadata
        num_trials: Number of trials to perform

    Returns
        best_hyperparameters : A tuple containing the optimal learning rate,
            weight decay, and scheduler type.
    """
    logger_tuning = Logger("hyperparameter_tuning", dataset=dataset_name, to_file=True)

    results = []

    for trial in range(num_trials):
        # Sample hyperparameters randomly
        learning_rate = random.choice(learning_rates)
        weight_decay = random.choice(weight_decays)
        scheduler_type = random.choice(scheduler_types)
        gamma_rate = random.choice(gamma_rates)

        gpf = GPF(dataset=dataset_name, groups=hierarchical_data)

        # Evaluate the performance with the sampled hyperparameters
        model, _ = gpf.train(
            lr=learning_rate,
            weight_decay=weight_decay,
            scheduler_type=scheduler_type,
            gamma_rate=gamma_rate,
        )

        val_loss = gpf.validate(model)

        results.append((learning_rate, weight_decay, scheduler_type, gamma_rate, val_loss))

    results.sort(key=lambda x: x[4])
    best_hyperparameters = results[0]

    logger_tuning.info(
        f"Best hyperparameters: learning rate = {best_hyperparameters[0]}, weight decay = {best_hyperparameters[1]}, scheduler = {best_hyperparameters[2]}, gamma = {best_hyperparameters[3]},"
    )
    logger_tuning.info(f"Validation loss: {best_hyperparameters[4]}")

    return best_hyperparameters
