from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.svm import SVC
from sklearn.neural_network import MLPClassifier
from sklearn.model_selection import GridSearchCV, StratifiedKFold

RANDOM_STATE = 42
CV_FOLDS     = 5


def get_models():
    """Returns {name: (estimator, param_grid)} for hyperparameter tuning."""
    return {
        "Logistic Regression": (
            LogisticRegression(max_iter=1000, random_state=RANDOM_STATE,
                               class_weight="balanced"),
            {"C": [0.01, 0.1, 1, 10],
             "solver": ["lbfgs", "liblinear"]},
        ),
        "Random Forest": (
            RandomForestClassifier(random_state=RANDOM_STATE,
                                   class_weight="balanced", n_jobs=-1),
            {"n_estimators": [100, 200],
             "max_depth": [None, 10, 20],
             "min_samples_split": [2, 5]},
        ),
        "SVM": (
            SVC(probability=True, random_state=RANDOM_STATE,
                class_weight="balanced"),
            {"C": [0.1, 1, 10],
             "kernel": ["rbf", "linear"],
             "gamma": ["scale"]},
        ),
        "Gradient Boosting": (
            GradientBoostingClassifier(random_state=RANDOM_STATE),
            {"n_estimators": [100, 200],
             "max_depth": [3, 5],
             "learning_rate": [0.05, 0.1],
             "subsample": [0.8, 1.0]},
        ),
        "ANN (MLP)": (
            MLPClassifier(max_iter=500, random_state=RANDOM_STATE,
                          early_stopping=True),
            {"hidden_layer_sizes": [(64,), (128, 64), (128, 64, 32)],
             "activation": ["relu", "tanh"],
             "alpha": [0.0001, 0.001]},
        ),
    }


def train_all_models(X_train, y_train, verbose=True):
    """
    Trains all 5 models via cross-validated grid search.
    Returns:
        fitted_models : {name: best_estimator}
        cv_scores     : {name: best_cv_AUC}
    """
    cv = StratifiedKFold(n_splits=CV_FOLDS, shuffle=True,
                         random_state=RANDOM_STATE)
    fitted_models, cv_scores = {}, {}

    for name, (model, grid) in get_models().items():
        if verbose:
            print(f"  ▸ Training {name}…", flush=True)
        search = GridSearchCV(
            model, grid, cv=cv, scoring="roc_auc",
            n_jobs=-1, refit=True, verbose=0
        )
        search.fit(X_train, y_train)
        fitted_models[name] = search.best_estimator_
        cv_scores[name]     = search.best_score_
        if verbose:
            print(f"      CV AUC = {search.best_score_:.4f} | "
                  f"best params: {search.best_params_}", flush=True)

    return fitted_models, cv_scores


if __name__ == "__main__":
    import numpy as np
    rng = np.random.default_rng(0)
    X = rng.standard_normal((200, 8))
    y = rng.integers(0, 2, 200)
    print("Quick smoke-test on random data…")
    fitted, cv = train_all_models(X, y)
    for n, s in cv.items():
        print(f"  {n:22s}: CV AUC = {s:.4f}")
    print("Models OK.")