import numpy as np
from xgboost import XGBClassifier
from sklearn.metrics import accuracy_score, f1_score, roc_auc_score
import sys
sys.path.insert(0, '.')
from data.data_loader import load_data, preprocess, split_and_scale
from src.smote import smote
from src.hyperparameter_tuning import grid_search

# run at import time so BEST_ALPHA/BEST_LAMBDA are available to both experiments
BEST_ALPHA, BEST_LAMBDA = grid_search()

XGB_PARAMS = dict(
    n_estimators=100, max_depth=4, learning_rate=0.05,
    subsample=0.7, colsample_bytree=0.7, min_child_weight=5,
    reg_alpha=BEST_ALPHA, reg_lambda=BEST_LAMBDA,
    eval_metric='logloss', early_stopping_rounds=15,
    random_state=42,
)


# searches val set for threshold maximising F1 — test set never touched
def find_best_threshold(y_true, proba):
    best_thresh, best_f1 = 0.5, 0.0
    for t in np.linspace(0.1, 0.9, 81):
        preds = (proba >= t).astype(int)
        f1    = f1_score(y_true, preds, zero_division=0)
        if f1 > best_f1:
            best_f1, best_thresh = f1, t
    return best_thresh, best_f1


def evaluate(model, X, y, name, threshold=0.5):
    proba  = model.predict_proba(X)[:, 1]
    preds  = (proba >= threshold).astype(int)
    acc    = accuracy_score(y, preds)
    f1     = f1_score(y, preds, zero_division=0)
    auc    = roc_auc_score(y, proba)
    suffix = f"  (thresh={threshold:.2f})" if threshold != 0.5 else ""
    print(f"  {name:48s} | F1: {f1:.3f} | AUC: {auc:.3f} | Acc: {acc:.3f}{suffix}")
    return acc, f1, auc


def run_condition(label, X_train, y_train, X_val, y_val, X_test, y_test,
                  use_smote=False, use_threshold_tuning=False):
    if use_smote:
        X_train, y_train = smote(X_train, np.array(y_train), k=5, random_state=42)

    model = XGBClassifier(**XGB_PARAMS)
    model.fit(X_train, y_train, eval_set=[(X_val, y_val)], verbose=False)

    threshold = 0.5
    if use_threshold_tuning:
        val_proba = model.predict_proba(X_val)[:, 1]
        threshold, _ = find_best_threshold(y_val, val_proba)

    _, f1, auc = evaluate(model, X_test, y_test, label, threshold=threshold)
    return f1, auc


def run_preprocessing_experiment(df):
    results = {}

    X, y, _ = preprocess(df.copy(), use_domain_cleaning=False)
    X_train, X_val, X_test, y_train, y_val, y_test, _ = split_and_scale(X, y)
    results['Baseline'] = run_condition(
        "Baseline", X_train, y_train, X_val, y_val, X_test, y_test)

    X, y, _ = preprocess(df.copy(), use_domain_cleaning=True)
    X_train, X_val, X_test, y_train, y_val, y_test, _ = split_and_scale(X, y)
    results['Domain cleaning'] = run_condition(
        "Domain cleaning", X_train, y_train, X_val, y_val, X_test, y_test)
    results['+ SMOTE'] = run_condition(
        "Domain cleaning + SMOTE",
        X_train, y_train, X_val, y_val, X_test, y_test, use_smote=True)
    results['+ Threshold'] = run_condition(
        "Domain cleaning + threshold tuning",
        X_train, y_train, X_val, y_val, X_test, y_test, use_threshold_tuning=True)
    results['Full pipeline'] = run_condition(
        "Full pipeline (cleaning + SMOTE + threshold)",
        X_train, y_train, X_val, y_val, X_test, y_test,
        use_smote=True, use_threshold_tuning=True)

    print(f"\n  {'Condition':<40} | {'F1':>6} | {'AUC':>6}")
    for name, (f1, auc) in results.items():
        marker = " ◀" if f1 == max(v[0] for v in results.values()) else ""
        print(f"  {name:<40} | {f1:>6.3f} | {auc:>6.3f}{marker}")


def regularization_experiment(df):
    from sklearn.metrics import roc_auc_score

    X, y, _ = preprocess(df.copy(), use_domain_cleaning=True)
    X_train, X_val, X_test, y_train, y_val, y_test, _ = split_and_scale(X, y)
    X_train_s, y_train_s = smote(X_train, np.array(y_train), random_state=42)

    def fit_eval(name, max_depth=8, min_child_weight=1,
                 reg_alpha=0, reg_lambda=0,
                 use_early_stopping=False, n_estimators=500):
        model = XGBClassifier(
            n_estimators=n_estimators, max_depth=max_depth,
            learning_rate=0.05, subsample=0.7, colsample_bytree=0.7,
            min_child_weight=min_child_weight,
            reg_alpha=reg_alpha, reg_lambda=reg_lambda,
            eval_metric='logloss', random_state=42,
            **(dict(early_stopping_rounds=20) if use_early_stopping else {}),
        )
        fit_kw = dict(eval_set=[(X_val, y_val)], verbose=False) \
                 if use_early_stopping else dict(verbose=False)
        model.fit(X_train_s, y_train_s, **fit_kw)

        stopped_at = (model.best_iteration + 1) if use_early_stopping else n_estimators
        train_auc  = roc_auc_score(y_train_s, model.predict_proba(X_train_s)[:, 1])
        test_auc   = roc_auc_score(y_test,    model.predict_proba(X_test)[:, 1])
        gap        = train_auc - test_auc
        print(f"  {name:<42s} | train: {train_auc:.3f} | test: {test_auc:.3f} | gap: {gap:.3f} | trees: {stopped_at}")
        return gap, test_auc

    print(f"\n  {'Condition':<42s} | train AUC | test AUC | gap   | trees")
    r = {}
    r['No reg']     = fit_eval("No regularization  (depth=8, min_child=1)")
    r['L2']         = fit_eval("L2 only  (reg_lambda=5.0)",          reg_lambda=5.0)
    r['L1']         = fit_eval("L1 only  (reg_alpha=2.0)",           reg_alpha=2.0)
    r['L1+L2']      = fit_eval("L1 + L2  (alpha=2.0, lambda=5.0)",  reg_alpha=2.0, reg_lambda=5.0)
    r['ES']         = fit_eval("Early stopping only  (depth=8)",     use_early_stopping=True)
    r['Production'] = fit_eval("L1 + L2 + early stopping  ◀ production",
                                max_depth=4, min_child_weight=5,
                                reg_alpha=BEST_ALPHA, reg_lambda=BEST_LAMBDA,
                                use_early_stopping=True)

    bg, ba = r['No reg']
    fg, fa = r['Production']
    print(f"\n  gap {bg:.3f} → {fg:.3f}  |  test AUC {ba:.3f} → {fa:.3f}")


if __name__ == '__main__':
    df = load_data()
    run_preprocessing_experiment(df)
    regularization_experiment(df)