"""Churn: compare model compositions on the same protocol as 05_Predictions/p04_churn.py.

Label = no purchase in the 8 weeks after the cut-off. 30% of outlets held out (same seed). Two validation folds
(train 2 Jun -> validate 16 Jun; train 2 + 16 Jun -> validate 30 Jun, held-out outlets) pick the composition; test =
train 2, 16, 30 Jun -> score 25 Aug on held-out outlets, once. Paired bootstrap gives the uncertainty of each gain.
"""
import sys, pickle, warnings, itertools
import numpy as np, pandas as pd
warnings.filterwarnings('ignore')
sys.path.insert(0, '/home/claude/binzagr/pred')
import p04_churn as C
from sklearn.metrics import roc_auc_score
from sklearn.ensemble import HistGradientBoostingClassifier, ExtraTreesClassifier, RandomForestClassifier
from catboost import CatBoostClassifier
from xgboost import XGBClassifier

rng = np.random.RandomState(20261023)
all_outlets = C.B.outlet_id.unique(); hold = set(rng.choice(all_outlets, size=int(0.3 * len(all_outlets)), replace=False))


def num_X(f):
    X = C.lin_X(f); X['channel_code'] = pd.factorize(f.channel)[0]; return X


def fit_predict(name, tr, te):
    if name == 'LightGBM': return C.model().fit(C.prep(tr), tr.y).predict_proba(C.prep(te))[:, 1]
    if name == 'Logistic regression': return C.logit().fit(C.lin_X(tr), tr.y).predict_proba(C.lin_X(te))[:, 1]
    if name == 'CatBoost':
        Xtr, Xte = tr[C.FEATS].copy(), te[C.FEATS].copy()
        m = CatBoostClassifier(iterations=600, learning_rate=0.03, depth=5, l2_leaf_reg=5, verbose=0, random_seed=7, thread_count=2)
        return m.fit(Xtr, tr.y, cat_features=['channel']).predict_proba(Xte)[:, 1]
    if name == 'XGBoost':
        Xtr, Xte = C.prep(tr), C.prep(te)
        m = XGBClassifier(n_estimators=400, learning_rate=0.03, max_depth=4, subsample=0.8, colsample_bytree=0.8, min_child_weight=5,
                          reg_lambda=2.0, tree_method='hist', enable_categorical=True, n_jobs=2, random_state=7)
        return m.fit(Xtr, tr.y).predict_proba(Xte)[:, 1]
    if name == 'HistGradientBoosting':
        m = HistGradientBoostingClassifier(max_iter=300, learning_rate=0.04, max_leaf_nodes=15, l2_regularization=1.0, random_state=7)
        return m.fit(num_X(tr), tr.y).predict_proba(num_X(te))[:, 1]
    if name == 'Random forest':
        m = RandomForestClassifier(n_estimators=500, min_samples_leaf=20, max_features=0.4, n_jobs=2, random_state=7)
        return m.fit(num_X(tr).fillna(-1), tr.y).predict_proba(num_X(te).fillna(-1))[:, 1]
    if name == 'Extra trees':
        m = ExtraTreesClassifier(n_estimators=500, min_samples_leaf=20, max_features=0.5, n_jobs=2, random_state=7)
        return m.fit(num_X(tr).fillna(-1), tr.y).predict_proba(num_X(te).fillna(-1))[:, 1]
    raise KeyError(name)


BASE = ['LightGBM', 'Logistic regression', 'CatBoost', 'XGBoost', 'HistGradientBoosting', 'Random forest', 'Extra trees']


def rank(v): return pd.Series(v).rank(pct=True).values


def all_scores(tr, te):
    sc = {m: fit_predict(m, tr, te) for m in BASE}
    for k in (2, 3):
        for combo in itertools.combinations(BASE, k):
            sc['Blend: ' + ' + '.join(combo)] = np.mean([rank(sc[m]) for m in combo], 0)
    sc['Blend: all seven'] = np.mean([rank(sc[m]) for m in BASE], 0)
    sc['Recency rule'] = te.recency_days.values.astype(float)
    return sc


if __name__ == '__main__':
    folds = [(['2026-06-02'], ['2026-06-16']), (['2026-06-02', '2026-06-16'], ['2026-06-30'])]
    val = {}
    for trc, vac in folds:
        tr = C.labelled(trc, outlets_drop=hold); va = C.labelled(vac, outlets_keep=hold)
        for k, v in all_scores(tr, va).items(): val.setdefault(k, []).append(roc_auc_score(va.y, v))
        print('fold', vac, 'done', flush=True)
    trt = C.labelled(['2026-06-02', '2026-06-16', '2026-06-30'], outlets_drop=hold)
    te = C.labelled(['2026-08-25'], outlets_keep=hold)
    st = all_scores(trt, te)
    rows = []
    for k, v in st.items():
        cap, prec = C.top_capture(te, v)
        rows.append({'model': k, 'val_auc_mean': np.mean(val[k]), 'val_auc_16jun': val[k][0], 'val_auc_30jun': val[k][1],
                     'test_auc': roc_auc_score(te.y, v), 'test_top20_capture': cap})
    R = pd.DataFrame(rows).sort_values('val_auc_mean', ascending=False)
    pick = R[R.model != 'Recency rule'].iloc[0].model
    # paired bootstrap of test AUC gains over LightGBM (the current model) and over the recency rule
    yb = te.y.values; B_ = 2000; rs = np.random.RandomState(1)
    def boot(a, b):
        d = []
        for _ in range(B_):
            i = rs.randint(0, len(yb), len(yb))
            if yb[i].min() == yb[i].max(): continue
            d.append(roc_auc_score(yb[i], st[a][i]) - roc_auc_score(yb[i], st[b][i]))
        return np.percentile(d, [5, 50, 95])
    ci = {m: boot(m, 'LightGBM') for m in [pick, 'Logistic regression', 'Blend: LightGBM + Logistic regression', 'CatBoost']}
    ci_rec = boot('LightGBM', 'Recency rule')
    print(R.head(15).round(4).to_string(index=False))
    print('picked on validation:', pick)
    print('test AUC gain vs LightGBM, 90% bootstrap interval:', {k: np.round(v, 4).tolist() for k, v in ci.items()})
    print('LightGBM gain vs recency rule, 90% interval:', np.round(ci_rec, 4).tolist())
    print('test outlets', len(te), 'went quiet', round(te.y.mean(), 4), 'positives', int(te.y.sum()))
    pickle.dump({'table': R, 'pick': pick, 'ci': ci, 'ci_rec': ci_rec, 'n_test': len(te), 'pos': int(te.y.sum()), 'rate': te.y.mean()},
                open('/home/claude/binzagr/boost/work/churn_race.pkl', 'wb'))