"""
Optuna hyperparameter tuning for XGBoost (GPU) and LightGBM (CPU).
"""

from __future__ import annotations

import numpy as np
import optuna
import torch
from sklearn.model_selection import StratifiedKFold, cross_val_score

optuna.logging.set_verbosity(optuna.logging.WARNING)

# GPU device detection (done once at import)
_XGB_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


# ═════════════════════════════════════════════════════════════════════════════
# XGBoost tuning (GPU)
# ═════════════════════════════════════════════════════════════════════════════

def tune_xgb(
    X: np.ndarray,
    y: np.ndarray,
    num_classes: int = 2,
    n_trials: int = 50,
    timeout: int = 300,
    cv: int = 3,
) -> dict:
    """Run Optuna search for XGBoost hyperparameters.

    Parameters
    ----------
    X : np.ndarray
        Feature matrix (float32).
    y : np.ndarray
        Labels.
    num_classes : int
        2 for binary, >2 for multi-class.
    n_trials : int
        Max Optuna trials.
    timeout : int
        Time limit in seconds.
    cv : int
        Cross-validation folds.

    Returns
    -------
    dict
        Best hyperparameters.
    """
    import xgboost as xgb

    objective = "binary:logistic" if num_classes == 2 else "multi:softprob"

    def _objective(trial):
        params = {
            "n_estimators": trial.suggest_int("n_estimators", 100, 800),
            "max_depth": trial.suggest_int("max_depth", 3, 10),
            "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.5, 1.0),
            "min_child_weight": trial.suggest_int("min_child_weight", 1, 10),
            "reg_alpha": trial.suggest_float("reg_alpha", 1e-5, 10.0, log=True),
            "reg_lambda": trial.suggest_float("reg_lambda", 1e-5, 10.0, log=True),
            "gamma": trial.suggest_float("gamma", 0.0, 5.0),
        }
        clf = xgb.XGBClassifier(
            **params,
            objective=objective,
            tree_method="hist",
            device=_XGB_DEVICE,
            eval_metric="logloss" if num_classes == 2 else "mlogloss",
            n_jobs=-1,
            random_state=42,
            verbosity=0,
        )
        skf = StratifiedKFold(n_splits=cv, shuffle=True, random_state=42)
        scores = cross_val_score(clf, X, y, cv=skf, scoring="accuracy", n_jobs=1)
        return scores.mean()

    study = optuna.create_study(direction="maximize")
    study.optimize(_objective, n_trials=n_trials, timeout=timeout, show_progress_bar=True)

    print(f"  Best XGBoost accuracy: {study.best_value:.4f}")
    print(f"  Best params: {study.best_params}")
    return study.best_params


# ═════════════════════════════════════════════════════════════════════════════
# LightGBM tuning (CPU)
# ═════════════════════════════════════════════════════════════════════════════

def tune_lgbm(
    X: np.ndarray,
    y: np.ndarray,
    num_classes: int = 2,
    n_trials: int = 50,
    timeout: int = 300,
    cv: int = 3,
) -> dict:
    """Run Optuna search for LightGBM hyperparameters.

    Parameters
    ----------
    X, y, num_classes, n_trials, timeout, cv : same as ``tune_xgb``.

    Returns
    -------
    dict
        Best hyperparameters.
    """
    import lightgbm as lgb

    objective = "binary" if num_classes == 2 else "multiclass"

    def _objective(trial):
        params = {
            "n_estimators": trial.suggest_int("n_estimators", 100, 800),
            "max_depth": trial.suggest_int("max_depth", 3, 12),
            "learning_rate": trial.suggest_float("learning_rate", 0.01, 0.3, log=True),
            "num_leaves": trial.suggest_int("num_leaves", 20, 150),
            "subsample": trial.suggest_float("subsample", 0.6, 1.0),
            "colsample_bytree": trial.suggest_float("colsample_bytree", 0.5, 1.0),
            "min_child_samples": trial.suggest_int("min_child_samples", 5, 50),
            "reg_alpha": trial.suggest_float("reg_alpha", 1e-5, 10.0, log=True),
            "reg_lambda": trial.suggest_float("reg_lambda", 1e-5, 10.0, log=True),
        }
        clf = lgb.LGBMClassifier(
            **params,
            objective=objective,
            class_weight="balanced",
            device_type="cpu",
            n_jobs=-1,
            random_state=42,
            verbose=-1,
        )
        skf = StratifiedKFold(n_splits=cv, shuffle=True, random_state=42)
        scores = cross_val_score(clf, X, y, cv=skf, scoring="accuracy", n_jobs=1)
        return scores.mean()

    study = optuna.create_study(direction="maximize")
    study.optimize(_objective, n_trials=n_trials, timeout=timeout, show_progress_bar=True)

    print(f"  Best LightGBM accuracy: {study.best_value:.4f}")
    print(f"  Best params: {study.best_params}")
    return study.best_params