"""
Adult Income Dataset - SOTA Solution
OpenML Task 7592 / data_id=1590
Target: AUC > 0.9300, Accuracy > 0.8756 on 10-fold CV
Method: LightGBM + XGBoost + CatBoost stacking + Feature Engineering + Optuna
"""

import warnings, sys
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
from sklearn.datasets import fetch_openml
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import roc_auc_score, accuracy_score
from sklearn.preprocessing import OrdinalEncoder
from sklearn.linear_model import LogisticRegression
import lightgbm as lgb
import xgboost as xgb
import catboost as cb
import optuna
optuna.logging.set_verbosity(optuna.logging.WARNING)
import time

def log(msg):
    print(msg, flush=True)
    sys.stdout.flush()

log("=" * 70)
log("ADULT INCOME DATASET - SOTA SOLUTION")
log("OpenML Task 7592 | Target: Acc > 0.8756, AUC > 0.9300")
log("=" * 70)

# ─────────────────────────────────────────────────────────
# 1. CHARGEMENT DONNÉES
# ─────────────────────────────────────────────────────────
log("\n[1/6] Chargement données OpenML (data_id=1590)...")
t0 = time.time()
X, y = fetch_openml(data_id=1590, as_frame=True, return_X_y=True, cache=True)
y_bin = (y == ">50K").astype(int)
log(f"  Shape: {X.shape} | Target: {y_bin.sum()} positifs / {len(y_bin)} total ({y_bin.mean():.1%})")

CAT_COLS = ["workclass", "education", "marital-status", "occupation",
            "relationship", "race", "sex", "native-country"]
NUM_COLS = ["age", "fnlwgt", "education-num", "capital-gain", "capital-loss", "hours-per-week"]

log("\n  EDA:")
for col in CAT_COLS:
    log(f"    {col:20s}: {X[col].nunique():3d} vals, {X[col].isna().sum():5d} NaN")
for col in NUM_COLS:
    log(f"    {col:20s}: mean={X[col].mean():.1f}, std={X[col].std():.1f}")
log(f"  Chargement: {time.time()-t0:.1f}s")

# ─────────────────────────────────────────────────────────
# 2. FEATURE ENGINEERING
# ─────────────────────────────────────────────────────────
log("\n[2/6] Feature Engineering avancé...")

CAT_COLS = ["workclass", "education", "marital-status", "occupation",
            "relationship", "race", "sex", "native-country"]
NUM_COLS = ["age", "fnlwgt", "education-num", "capital-gain", "capital-loss", "hours-per-week"]

def build_features(X, fit_encoder=True, encoder=None):
    age      = X["age"].astype(float).values
    fnlwgt   = X["fnlwgt"].astype(float).values
    edu_num  = X["education-num"].astype(float).values
    cap_gain = X["capital-gain"].astype(float).values
    cap_loss = X["capital-loss"].astype(float).values
    hours    = X["hours-per-week"].astype(float).values

    X_num = np.column_stack([
        age, fnlwgt, edu_num, cap_gain, cap_loss, hours,
        np.log1p(cap_gain), np.log1p(cap_loss),
        cap_gain - cap_loss,
        np.log1p(np.abs(cap_gain - cap_loss)) * np.sign(cap_gain - cap_loss),
        ((cap_gain > 0) | (cap_loss > 0)).astype(float),
        (cap_gain > 0).astype(float), (cap_loss > 0).astype(float),
        age ** 2,
        pd.cut(age, bins=[0,25,35,45,55,65,100], labels=False).astype(float),
        pd.cut(hours, bins=[0,35,40,45,60,100], labels=False).astype(float),
        (hours > 40).astype(float),
        np.log1p(fnlwgt),
        edu_num * age, edu_num * hours
    ])

    X_cat = X[CAT_COLS].astype(str)
    if fit_encoder:
        encoder = OrdinalEncoder(handle_unknown="use_encoded_value", unknown_value=-1)
        X_cat_enc = encoder.fit_transform(X_cat)
    else:
        X_cat_enc = encoder.transform(X_cat)

    return np.hstack([X_num, X_cat_enc]), encoder

def build_cb_features(X):
    X_cb = X.copy()
    for col in CAT_COLS:
        if hasattr(X_cb[col], 'cat'):
            X_cb[col] = X_cb[col].cat.add_categories(["Unknown"]).fillna("Unknown").astype(str)
        else:
            X_cb[col] = X_cb[col].fillna("Unknown").astype(str)
    cap_gain = X_cb["capital-gain"].astype(float)
    cap_loss = X_cb["capital-loss"].astype(float)
    X_cb["cap_gain_log"] = np.log1p(cap_gain)
    X_cb["cap_loss_log"] = np.log1p(cap_loss)
    X_cb["cap_net"] = cap_gain - cap_loss
    X_cb["cap_any"] = ((cap_gain > 0) | (cap_loss > 0)).astype(float)
    X_cb["age_bins"] = pd.cut(X_cb["age"].astype(float), bins=[0,25,35,45,55,65,100], labels=False).astype(float)
    X_cb["edu_x_age"] = X_cb["education-num"].astype(float) * X_cb["age"].astype(float)
    X_cb["fnlwgt_log"] = np.log1p(X_cb["fnlwgt"].astype(float))
    return X_cb

X_enc, oe = build_features(X)
X_cb_df    = build_cb_features(X)
y_arr      = y_bin.values
n          = len(y_arr)

log(f"  Features LGB/XGB: {X_enc.shape[1]} | CatBoost: {X_cb_df.shape[1]} colonnes")

# ─────────────────────────────────────────────────────────
# 3. BASELINE 3-FOLD (validation rapide architecture)
# ─────────────────────────────────────────────────────────
log("\n[3/6] Baseline 3-fold CV (300 estimators, validation architecture)...")

# Paramètres baseline réduits pour vitesse
LGB_BASE = dict(n_estimators=300, learning_rate=0.05, num_leaves=63,
                colsample_bytree=0.8, subsample=0.8, subsample_freq=1,
                min_child_samples=20, reg_alpha=0.05, reg_lambda=1.0,
                max_depth=8, random_state=42, n_jobs=-1, verbose=-1)
XGB_BASE = dict(n_estimators=300, learning_rate=0.05, max_depth=6,
                colsample_bytree=0.8, subsample=0.8, min_child_weight=5,
                reg_alpha=0.05, reg_lambda=1.5, eval_metric="logloss",
                random_state=42, n_jobs=-1, verbosity=0)
CB_BASE  = dict(iterations=300, learning_rate=0.05, depth=8,
                cat_features=CAT_COLS, random_seed=42, verbose=0, thread_count=4)

cv3 = StratifiedKFold(n_splits=3, shuffle=True, random_state=42)
oof_lgb_3 = np.zeros(n); oof_xgb_3 = np.zeros(n); oof_cb_3 = np.zeros(n)

for fi, (tr, va) in enumerate(cv3.split(X_enc, y_arr)):
    t_fold = time.time()
    log(f"  Fold {fi+1}/3 LGB...", )
    m = lgb.LGBMClassifier(**LGB_BASE); m.fit(X_enc[tr], y_arr[tr])
    oof_lgb_3[va] = m.predict_proba(X_enc[va])[:, 1]
    log(f"  Fold {fi+1}/3 XGB...")
    m = xgb.XGBClassifier(**XGB_BASE); m.fit(X_enc[tr], y_arr[tr])
    oof_xgb_3[va] = m.predict_proba(X_enc[va])[:, 1]
    log(f"  Fold {fi+1}/3 CB ...")
    m = cb.CatBoostClassifier(**CB_BASE); m.fit(X_cb_df.iloc[tr], y_arr[tr])
    oof_cb_3[va] = m.predict_proba(X_cb_df.iloc[va])[:, 1]
    avg = (oof_lgb_3[va] + oof_xgb_3[va] + oof_cb_3[va]) / 3
    log(f"  → Fold {fi+1} done: AUC={roc_auc_score(y_arr[va], avg):.5f} Acc={accuracy_score(y_arr[va], (avg>=0.5).astype(int)):.5f} ({time.time()-t_fold:.0f}s)")

avg_3 = (oof_lgb_3 + oof_xgb_3 + oof_cb_3) / 3
auc_avg_3 = roc_auc_score(y_arr, avg_3)
best_acc_3 = max(accuracy_score(y_arr, (avg_3 >= t).astype(int)) for t in np.arange(0.3, 0.7, 0.005))
log(f"\n  BASELINE 3-FOLD: LGB={roc_auc_score(y_arr, oof_lgb_3):.5f} "
    f"XGB={roc_auc_score(y_arr, oof_xgb_3):.5f} CB={roc_auc_score(y_arr, oof_cb_3):.5f} "
    f"AVG_AUC={auc_avg_3:.5f} BestAcc={best_acc_3:.5f}")
log(f"  Target 0.8756: {'✅ ATTEINT' if best_acc_3 >= 0.8756 else '❌ ' + str(round(best_acc_3,5))}")

# ─────────────────────────────────────────────────────────
# 4. OPTUNA TUNING
# ─────────────────────────────────────────────────────────
log("\n[4/6] Optuna Tuning...")
cv_inner = StratifiedKFold(n_splits=3, shuffle=True, random_state=123)

# LightGBM - 40 trials
log("  Tuning LightGBM (40 trials)...")
def lgb_obj(trial):
    p = dict(
        n_estimators    = trial.suggest_int("n_estimators", 200, 1200),
        learning_rate   = trial.suggest_float("learning_rate", 0.01, 0.1, log=True),
        num_leaves      = trial.suggest_int("num_leaves", 31, 127),
        max_depth       = trial.suggest_int("max_depth", 4, 10),
        min_child_samples = trial.suggest_int("min_child_samples", 5, 80),
        colsample_bytree= trial.suggest_float("colsample_bytree", 0.5, 1.0),
        subsample       = trial.suggest_float("subsample", 0.5, 1.0),
        subsample_freq  = 1,
        reg_alpha       = trial.suggest_float("reg_alpha", 1e-4, 5.0, log=True),
        reg_lambda      = trial.suggest_float("reg_lambda", 1e-4, 5.0, log=True),
        random_state=42, n_jobs=-1, verbose=-1
    )
    return np.mean([roc_auc_score(y_arr[va],
        lgb.LGBMClassifier(**p).fit(X_enc[tr], y_arr[tr]).predict_proba(X_enc[va])[:,1])
        for tr, va in cv_inner.split(X_enc, y_arr)])

st_lgb = optuna.create_study(direction="maximize", sampler=optuna.samplers.TPESampler(seed=42))
st_lgb.optimize(lgb_obj, n_trials=40, show_progress_bar=False)
best_lgb = st_lgb.best_params
log(f"  LGB best AUC={st_lgb.best_value:.5f} | {best_lgb}")

# XGBoost - 40 trials
log("  Tuning XGBoost (40 trials)...")
def xgb_obj(trial):
    p = dict(
        n_estimators     = trial.suggest_int("n_estimators", 200, 1200),
        learning_rate    = trial.suggest_float("learning_rate", 0.01, 0.1, log=True),
        max_depth        = trial.suggest_int("max_depth", 3, 10),
        min_child_weight = trial.suggest_int("min_child_weight", 1, 20),
        colsample_bytree = trial.suggest_float("colsample_bytree", 0.5, 1.0),
        subsample        = trial.suggest_float("subsample", 0.5, 1.0),
        gamma            = trial.suggest_float("gamma", 0, 3),
        reg_alpha        = trial.suggest_float("reg_alpha", 1e-4, 5.0, log=True),
        reg_lambda       = trial.suggest_float("reg_lambda", 1e-4, 5.0, log=True),
        eval_metric="logloss", random_state=42, n_jobs=-1, verbosity=0
    )
    return np.mean([roc_auc_score(y_arr[va],
        xgb.XGBClassifier(**p).fit(X_enc[tr], y_arr[tr]).predict_proba(X_enc[va])[:,1])
        for tr, va in cv_inner.split(X_enc, y_arr)])

st_xgb = optuna.create_study(direction="maximize", sampler=optuna.samplers.TPESampler(seed=42))
st_xgb.optimize(xgb_obj, n_trials=40, show_progress_bar=False)
best_xgb = st_xgb.best_params
log(f"  XGB best AUC={st_xgb.best_value:.5f} | {best_xgb}")

# CatBoost - 25 trials (plus lent)
log("  Tuning CatBoost (25 trials)...")
def cb_obj(trial):
    p = dict(
        iterations        = trial.suggest_int("iterations", 200, 800),
        learning_rate     = trial.suggest_float("learning_rate", 0.01, 0.1, log=True),
        depth             = trial.suggest_int("depth", 4, 9),
        l2_leaf_reg       = trial.suggest_float("l2_leaf_reg", 0.01, 10.0, log=True),
        bagging_temperature = trial.suggest_float("bagging_temperature", 0, 3),
        random_strength   = trial.suggest_float("random_strength", 0, 3),
        cat_features=CAT_COLS, random_seed=42, verbose=0, thread_count=4
    )
    return np.mean([roc_auc_score(y_arr[va],
        cb.CatBoostClassifier(**p).fit(X_cb_df.iloc[tr], y_arr[tr]).predict_proba(X_cb_df.iloc[va])[:,1])
        for tr, va in cv_inner.split(X_enc, y_arr)])

st_cb = optuna.create_study(direction="maximize", sampler=optuna.samplers.TPESampler(seed=42))
st_cb.optimize(cb_obj, n_trials=25, show_progress_bar=False)
best_cb = st_cb.best_params
log(f"  CB  best AUC={st_cb.best_value:.5f} | {best_cb}")

# ─────────────────────────────────────────────────────────
# 5. STACKING FINAL 10-FOLD
# ─────────────────────────────────────────────────────────
log("\n[5/6] Stacking Final 10-Fold CV (paramètres Optuna)...")

# Paramètres finaux tunés
lgb_final = {**best_lgb, "random_state": 42, "n_jobs": -1, "verbose": -1}
xgb_final = {**best_xgb, "eval_metric": "logloss", "random_state": 42, "n_jobs": -1, "verbosity": 0}
cb_final  = {**best_cb, "cat_features": CAT_COLS, "random_seed": 42, "verbose": 0, "thread_count": 4}

cv10 = StratifiedKFold(n_splits=10, shuffle=True, random_state=42)
oof_lgb = np.zeros(n); oof_xgb = np.zeros(n); oof_cb = np.zeros(n)
fold_aucs = []

for fi, (tr, va) in enumerate(cv10.split(X_enc, y_arr)):
    t_f = time.time()
    log(f"  Fold {fi+1:2d}/10 LGB...")
    m_lgb = lgb.LGBMClassifier(**lgb_final); m_lgb.fit(X_enc[tr], y_arr[tr])
    oof_lgb[va] = m_lgb.predict_proba(X_enc[va])[:, 1]
    log(f"  Fold {fi+1:2d}/10 XGB...")
    m_xgb = xgb.XGBClassifier(**xgb_final); m_xgb.fit(X_enc[tr], y_arr[tr])
    oof_xgb[va] = m_xgb.predict_proba(X_enc[va])[:, 1]
    log(f"  Fold {fi+1:2d}/10 CB ...")
    m_cb = cb.CatBoostClassifier(**cb_final); m_cb.fit(X_cb_df.iloc[tr], y_arr[tr])
    oof_cb[va] = m_cb.predict_proba(X_cb_df.iloc[va])[:, 1]
    avg = (oof_lgb[va] + oof_xgb[va] + oof_cb[va]) / 3
    fold_auc = roc_auc_score(y_arr[va], avg)
    fold_aucs.append(fold_auc)
    log(f"  → Fold {fi+1:2d} done: AUC={fold_auc:.5f} ({time.time()-t_f:.0f}s)")

# ─────────────────────────────────────────────────────────
# 6. RÉSULTATS + META-STACKING
# ─────────────────────────────────────────────────────────
log("\n[6/6] Résultats finaux...")

auc_lgb = roc_auc_score(y_arr, oof_lgb)
auc_xgb = roc_auc_score(y_arr, oof_xgb)
auc_cb  = roc_auc_score(y_arr, oof_cb)

# Moyenne simple + threshold sweep
avg = (oof_lgb + oof_xgb + oof_cb) / 3
auc_avg = roc_auc_score(y_arr, avg)
acc_05  = accuracy_score(y_arr, (avg >= 0.5).astype(int))
best_acc_avg, best_thr_avg = max(
    ((accuracy_score(y_arr, (avg >= t).astype(int)), t) for t in np.arange(0.3, 0.70, 0.002)),
    key=lambda x: x[0])

# Weighted blend grid search
best_auc_w, best_w = 0, (1/3, 1/3, 1/3)
for w1 in np.arange(0.1, 0.7, 0.1):
    for w2 in np.arange(0.1, 0.7, 0.1):
        w3 = 1.0 - w1 - w2
        if w3 <= 0.05: continue
        auc = roc_auc_score(y_arr, w1*oof_lgb + w2*oof_xgb + w3*oof_cb)
        if auc > best_auc_w:
            best_auc_w, best_w = auc, (w1, w2, w3)

wblend = best_w[0]*oof_lgb + best_w[1]*oof_xgb + best_w[2]*oof_cb
best_acc_w = max(accuracy_score(y_arr, (wblend >= t).astype(int)) for t in np.arange(0.3, 0.70, 0.002))

# Meta-stacking LogReg
log("  Meta-stacking LogReg...")
meta_X = np.column_stack([oof_lgb, oof_xgb, oof_cb])
oof_meta = np.zeros(n)
for tr, va in cv10.split(meta_X, y_arr):
    lr = LogisticRegression(C=10, max_iter=1000, random_state=42)
    lr.fit(meta_X[tr], y_arr[tr])
    oof_meta[va] = lr.predict_proba(meta_X[va])[:, 1]
auc_meta = roc_auc_score(y_arr, oof_meta)
best_acc_meta = max(accuracy_score(y_arr, (oof_meta >= t).astype(int)) for t in np.arange(0.3, 0.70, 0.002))

# Meilleurs scores finaux
best_auc_all = max(auc_avg, best_auc_w, auc_meta)
best_acc_all = max(best_acc_avg, best_acc_w, best_acc_meta)

# ─────────────────────────────────────────────────────────
# RAPPORT
# ─────────────────────────────────────────────────────────
log("\n" + "=" * 70)
log("RAPPORT FINAL - ADULT INCOME DATASET")
log("=" * 70)
log("\n📊 RÉSULTATS 10-FOLD CV:")
log(f"  LightGBM  seul  - AUC: {auc_lgb:.5f}")
log(f"  XGBoost   seul  - AUC: {auc_xgb:.5f}")
log(f"  CatBoost  seul  - AUC: {auc_cb:.5f}")
log(f"  Moyenne simple  - AUC: {auc_avg:.5f} | Acc@0.5={acc_05:.5f} | Acc@opt={best_acc_avg:.5f} (thr={best_thr_avg:.3f})")
log(f"  Poids optim     - AUC: {best_auc_w:.5f} | Acc@opt={best_acc_w:.5f} (w={best_w[0]:.1f}/{best_w[1]:.1f}/{best_w[2]:.1f})")
log(f"  Meta-LR stack   - AUC: {auc_meta:.5f} | Acc@opt={best_acc_meta:.5f}")
log(f"\n  AUC fold-by-fold: {[round(x,4) for x in fold_aucs]}")
log(f"  Mean±Std: {np.mean(fold_aucs):.5f} ± {np.std(fold_aucs):.5f}")
log(f"\n🏆 MEILLEURE: AUC={best_auc_all:.5f} | Acc={best_acc_all:.5f}")
log(f"\n🎯 OBJECTIFS:")
log(f"   Accuracy > 0.8756: {'✅ ATTEINT (' + str(round(best_acc_all,5)) + ')' if best_acc_all > 0.8756 else '❌ ' + str(round(best_acc_all,5))}")
log(f"   AUC      > 0.9300: {'✅ ATTEINT (' + str(round(best_auc_all,5)) + ')' if best_auc_all > 0.9300 else '❌ ' + str(round(best_auc_all,5))}")
log(f"\n📋 vs OpenML SOTA (AdaBoost 2017: AUC=0.92840 Acc=0.87400):")
log(f"   ΔAUC: {best_auc_all - 0.92840:+.5f} | ΔAcc: {best_acc_all - 0.87400:+.5f}")
log("\n" + "=" * 70)
log("TERMINÉ.")
log("=" * 70)