"""
automl_lib.tuner
~~~~~~~~~~~~~~~~
Optuna-based hyperparameter optimization.
"""

from __future__ import annotations

from typing import Any, Callable

import optuna

from .config import AutoMLConfig
from .utils import get_logger

logger = get_logger(__name__)

# Silence Optuna's verbose logging
optuna.logging.set_verbosity(optuna.logging.WARNING)


def tune_hyperparameters(
    objective_fn: Callable[[optuna.Trial], float],
    config: AutoMLConfig,
    direction: str = "maximize",
) -> dict[str, Any]:
    """Run Optuna HPO and return the best trial's params.

    Parameters
    ----------
    objective_fn : callable
        A function ``f(trial) -> metric_value`` that builds, trains, and
        evaluates a model using trial's suggestions.
    config : AutoMLConfig
        Controls ``tune_trials`` and ``tune_timeout``.
    direction : str
        ``"maximize"`` for accuracy-like metrics or ``"minimize"`` for loss.

    Returns
    -------
    dict
        Best hyperparameters found.
    """
    study = optuna.create_study(direction=direction)
    study.optimize(
        objective_fn,
        n_trials=config.tune_trials,
        timeout=config.tune_timeout,
        show_progress_bar=True,
    )

    best = study.best_trial
    logger.info(f"Best trial #{best.number}: value={best.value:.4f}")
    logger.info(f"   Params: {best.params}")
    return best.params


def suggest_params(trial: optuna.Trial) -> dict[str, Any]:
    """Suggest a standard set of hyperparameters for neural-network training.

    Use this inside your objective function:

    .. code-block:: python

        def objective(trial):
            params = suggest_params(trial)
            # build & train model with params …
    """
    return {
        "learning_rate": trial.suggest_float("learning_rate", 1e-5, 1e-1, log=True),
        "batch_size": trial.suggest_categorical("batch_size", [16, 32, 64, 128]),
        "optimizer": trial.suggest_categorical("optimizer", ["adam", "adamw", "sgd"]),
        "dropout": trial.suggest_float("dropout", 0.1, 0.5),
        "weight_decay": trial.suggest_float("weight_decay", 1e-6, 1e-2, log=True),
        "hidden_1": trial.suggest_int("hidden_1", 64, 512, step=64),
        "hidden_2": trial.suggest_int("hidden_2", 32, 256, step=32),
    }


def suggest_sklearn_params(trial: optuna.Trial, model_type: str) -> dict[str, Any]:
    """Suggest hyperparameters for scikit-learn / XGBoost models."""
    if model_type in ("sklearn_rf", "random_forest"):
        return {
            "n_estimators": trial.suggest_int("n_estimators", 50, 500, step=50),
            "max_depth": trial.suggest_int("max_depth", 3, 20),
            "min_samples_split": trial.suggest_int("min_samples_split", 2, 20),
            "min_samples_leaf": trial.suggest_int("min_samples_leaf", 1, 10),
        }
    if model_type in ("sklearn_xgb", "xgboost"):
        return {
            "n_estimators": trial.suggest_int("n_estimators", 50, 500, step=50),
            "max_depth": trial.suggest_int("max_depth", 3, 12),
            "learning_rate": trial.suggest_float("learning_rate", 0.01, 0.3, log=True),
            "subsample": trial.suggest_float("subsample", 0.6, 1.0),
            "colsample_bytree": trial.suggest_float("colsample_bytree", 0.6, 1.0),
        }
    return {}
