import numpy as np
import xgboost as xgb
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import roc_auc_score

def evaluate_max_depth(X, y, max_depths=[3, 5, 7, 10]):
    results = {}

    for depth in max_depths:
        print(f"Evaluating max_depth = {depth}")
        skf = StratifiedKFold(n_splits=3, shuffle=True, random_state=42)
        scores = []

        for train_idx, val_idx in skf.split(X, y):
            model = xgb.XGBClassifier(
                max_depth=depth,
                tree_method="gpu_hist",
                predictor="gpu_predictor",
                use_label_encoder=False,
                eval_metric="auc"
            )

            model.fit(X[train_idx], y[train_idx])
            preds = model.predict_proba(X[val_idx])[:, 1]
            scores.append(roc_auc_score(y[val_idx], preds))

        avg_score = np.mean(scores)
        std_score = np.std(scores)
        results[depth] = {'mean': avg_score, 'std': std_score}

        print(f"Max depth {depth}: {avg_score:.4f} ± {std_score:.4f}")

    return results

        __
        __