"""

automl_lib.tuner

~~~~~~~~~~~~~~~~

Optuna-based hyperparameter optimization.
"""

from __future__ import annotations


from typing import Any, Callable

import optuna

from .config import trainiqConfig
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: trainiqConfig,

    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 : trainiqConfig

        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 {}

