"""Outlet churn: will an active outlet go quiet (no purchase) in the next 8 weeks?

Back-test: trained on cut-offs 5 May, 2 Jun and 30 Jun (labels end 25 Aug); tested at the 25 Aug cut-off on outlets the
model never saw (30% of outlets held out), against the recency rule (days since last purchase). Final model: all outlets,
cut-offs 2 Jun - 25 Aug; scores every outlet active in the 13 weeks to 20 Oct for the 8 weeks after it.
"""
import numpy as np, pandas as pd, lightgbm as lgb, pickle
from sklearn.metrics import roc_auc_score
from pcommon import *

S = pd.read_parquet(f'{WORK}/sales.parquet'); R = pd.read_parquet(f'{WORK}/returns.parquet')
C = pd.read_parquet(f'{CLEAN}/dim_customer.parquet'); C['outlet_id'] = outlet_id(C.dist_cd, C.cust_cd)
C = C.set_index('outlet_id')
B = S[S.sellable & (S.net_amt > 0)]
H = pd.Timedelta(days=56)
P = pd.read_parquet(f'{WORK}/products.parquet').set_index('sku')
B = B.assign(principal=B.sku.map(P.l1_code))


def features(cut):
    past = B[B.date <= cut]
    w13 = past[past.date > cut - pd.Timedelta(days=91)]
    outs = w13.outlet_id.unique()
    g = past[past.outlet_id.isin(outs)]
    f = pd.DataFrame(index=pd.Index(outs, name='outlet_id'))
    last = g.groupby('outlet_id').date.max(); first = g.groupby('outlet_id').date.min()
    f['recency_days'] = (cut - last).dt.days
    f['tenure_days'] = (cut - first).dt.days
    for wk in (4, 8, 13, 26):
        x = g[g.date > cut - pd.Timedelta(weeks=wk)]
        f[f'days_{wk}w'] = x.groupby('outlet_id').date.nunique()
        f[f'sales_{wk}w'] = x.groupby('outlet_id').net_amt.sum()
    f = f.fillna(0)
    prev8 = g[(g.date <= cut - pd.Timedelta(weeks=8)) & (g.date > cut - pd.Timedelta(weeks=16))].groupby('outlet_id').net_amt.sum()
    f['trend_8w'] = (f.sales_8w + 1) / (f.index.map(prev8).fillna(0) + 1)
    x13 = g[g.date > cut - pd.Timedelta(weeks=13)]
    f['invoices_13w'] = x13.groupby('outlet_id').doc_no.nunique()
    f['avg_invoice_13w'] = f.sales_13w / f.invoices_13w.clip(lower=1)
    f['skus_13w'] = x13.groupby('outlet_id').sku.nunique(); f['principals_13w'] = x13.groupby('outlet_id').principal.nunique()
    f['vans_13w'] = x13.groupby('outlet_id').van.nunique()
    days = g[g.date > cut - pd.Timedelta(weeks=13)].groupby('outlet_id').date.apply(lambda s: np.diff(np.sort(s.unique()).astype('datetime64[D]').astype(int)))
    f['median_gap'] = days.map(lambda d: np.median(d) if len(d) else np.nan); f['gap_cv'] = days.map(lambda d: np.std(d) / np.mean(d) if len(d) > 1 else np.nan)
    f['recency_vs_gap'] = f.recency_days / f.median_gap.fillna(91)
    rr = R[(R.date <= cut) & (R.date > cut - pd.Timedelta(weeks=13))].groupby('outlet_id').net_amt.sum()
    f['return_rate_13w'] = f.index.map(rr).fillna(0) / f.sales_13w.clip(lower=1)
    allx = S[(S.date <= cut) & (S.date > cut - pd.Timedelta(weeks=13))]
    foc = allx.groupby('outlet_id').apply(lambda d: d.loc[d.is_foc, 'qty_cases'].sum() / max(d.qty_cases.sum(), 1e-9))
    f['foc_share_13w'] = f.index.map(foc).fillna(0)
    # the outlet's main van: is it still active, and how is it trending?
    mv = x13.groupby('outlet_id').van.agg(lambda s: s.value_counts().index[0]); f['main_van_share'] = x13.groupby('outlet_id').van.agg(lambda s: s.value_counts(normalize=True).iloc[0])
    vact = past[past.date > cut - pd.Timedelta(days=14)].groupby('van').date.nunique()
    vtr = (past[past.date > cut - pd.Timedelta(weeks=4)].groupby('van').net_amt.sum() + 1) / (past[(past.date <= cut - pd.Timedelta(weeks=4)) & (past.date > cut - pd.Timedelta(weeks=8))].groupby('van').net_amt.sum() + 1)
    f['main_van_days_14d'] = mv.map(vact).fillna(0).reindex(f.index).values; f['main_van_trend'] = mv.map(vtr).fillna(0).reindex(f.index).values
    f['main_van'] = mv.reindex(f.index).values
    cc = C.reindex(f.index)
    f['credit'] = (cc.cust_type == 'Q').astype(int).values; f['limit_placeholder'] = cc.flag_credit_limit_sentinel.fillna(False).astype(int).values
    f['channel'] = cc.cust_hier2.fillna('NA').astype(str).values; f['lat'] = cc.latitude.values; f['lon'] = cc.longitude.values
    f['account_age_days'] = (cut - cc.cust_opendt).dt.days.values
    f['cut'] = cut
    return f.reset_index()


def label(f, cut):
    fut = set(B[(B.date > cut) & (B.date <= cut + H)].outlet_id)
    return (~f.outlet_id.isin(fut)).astype(int)


FEATS = ['recency_days', 'days_4w', 'sales_4w', 'days_8w', 'sales_8w', 'days_13w', 'sales_13w', 'trend_8w',
         'invoices_13w', 'avg_invoice_13w', 'skus_13w', 'principals_13w', 'vans_13w', 'median_gap', 'gap_cv', 'recency_vs_gap', 'return_rate_13w',
         'foc_share_13w', 'main_van_share', 'main_van_days_14d', 'main_van_trend', 'credit', 'limit_placeholder', 'channel']


def prep(f):
    X = f[FEATS].copy(); X['channel'] = pd.Categorical(X.channel, categories=sorted(B_channels)); return X


def model():
    return lgb.LGBMClassifier(n_estimators=300, learning_rate=0.03, num_leaves=15, min_child_samples=40, subsample=0.8, subsample_freq=1,
                              colsample_bytree=0.8, reg_lambda=2.0, verbose=-1, n_jobs=2)


B_channels = set(C.cust_hier2.fillna('NA').astype(str)) | {'NA'}

from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.calibration import CalibratedClassifierCV
NUM = [c for c in FEATS if c != 'channel']


def logit():
    return make_pipeline(StandardScaler(), LogisticRegression(C=0.5, max_iter=2000))


def lin_X(f):
    X = f[NUM].copy().fillna(f[NUM].median(numeric_only=True))
    for c in ('sales_4w', 'sales_8w', 'sales_13w', 'avg_invoice_13w', 'recency_days', 'median_gap', 'trend_8w', 'recency_vs_gap', 'main_van_trend'):
        X[c] = np.log1p(X[c].clip(lower=0))
    return X


def top_capture(d, score, q=0.2):
    n = int(np.ceil(q * len(d))); top = d.assign(s=np.asarray(score)).sort_values('s', ascending=False).head(n)
    return top.y.sum() / d.y.sum(), top.y.mean()


def labelled(cuts, outlets_keep=None, outlets_drop=None):
    out = []
    for c in cuts:
        f = features(pd.Timestamp(c)); f['y'] = label(f, pd.Timestamp(c)); out.append(f)
    out = pd.concat(out, ignore_index=True)
    if outlets_keep is not None: out = out[out.outlet_id.isin(outlets_keep)]
    if outlets_drop is not None: out = out[~out.outlet_id.isin(outlets_drop)]
    return out


def score_variants(tr, te):
    r = {}
    m = model().fit(prep(tr), tr.y); r['lightgbm'] = m.predict_proba(prep(te))[:, 1]
    lm = logit().fit(lin_X(tr), tr.y); r['logistic'] = lm.predict_proba(lin_X(te))[:, 1]
    r['blend'] = pd.Series(r['lightgbm']).rank(pct=True).values * 0.5 + pd.Series(r['logistic']).rank(pct=True).values * 0.5
    r['recency_rule'] = te.recency_days.values
    return r


if __name__ == '__main__':
    rng = np.random.RandomState(20261023)
    all_outlets = B.outlet_id.unique(); hold = set(rng.choice(all_outlets, size=int(0.3 * len(all_outlets)), replace=False))
    # 1) choose the variant on validation: train cut-off 2 Jun, validate at 30 Jun on held-out outlets
    trv = labelled(['2026-06-02', '2026-06-16'], outlets_drop=hold); vav = labelled(['2026-06-30'], outlets_keep=hold)
    sv = score_variants(trv, vav)
    val_auc = {k: roc_auc_score(vav.y, v) for k, v in sv.items()}
    print('validation AUC', {k: round(v, 4) for k, v in val_auc.items()})
    choice = max([k for k in val_auc if k != 'recency_rule'], key=val_auc.get)
    # 2) test: train cut-offs 2, 16, 30 Jun (labels end 25 Aug), test at 25 Aug on held-out outlets
    trt = labelled(['2026-06-02', '2026-06-16', '2026-06-30'], outlets_drop=hold); te = labelled([str(RECO_CUTOFF.date())], outlets_keep=hold)
    st = score_variants(trt, te)
    res = {'chosen_variant': choice, 'test_outlets': len(te), 'went_quiet_rate': te.y.mean()}
    for k, v in st.items():
        res[f'auc_{k}'] = roc_auc_score(te.y, v); res[f'top20_capture_{k}'], res[f'top20_precision_{k}'] = top_capture(te, v)
    for k, v in val_auc.items(): res[f'validation_auc_{k}'] = v
    te = te.assign(p=st[choice])
    print({k: (round(float(v), 4) if not isinstance(v, str) else v) for k, v in res.items()})
    # 3) final model: all outlets, cut-offs 2 Jun - 25 Aug (labels complete), isotonic-calibrated probabilities
    fin = labelled(['2026-06-02', '2026-06-16', '2026-06-30', '2026-07-14', '2026-07-28', '2026-08-11', '2026-08-25'])
    now = features(DATA_END)
    if choice == 'logistic':
        mf = CalibratedClassifierCV(logit(), method='isotonic', cv=5).fit(lin_X(fin), fin.y); now['p'] = mf.predict_proba(lin_X(now))[:, 1]
    elif choice == 'lightgbm':
        mf = CalibratedClassifierCV(model(), method='isotonic', cv=5).fit(prep(fin), fin.y); now['p'] = mf.predict_proba(prep(now))[:, 1]
    else:
        a = CalibratedClassifierCV(model(), method='isotonic', cv=5).fit(prep(fin), fin.y).predict_proba(prep(now))[:, 1]
        b = CalibratedClassifierCV(logit(), method='isotonic', cv=5).fit(lin_X(fin), fin.y).predict_proba(lin_X(now))[:, 1]
        now['p'] = (a + b) / 2
    dec = te.assign(b=pd.qcut(te.p.rank(method='first'), 5, labels=False)).groupby('b').agg(score=('p', 'mean'), went_quiet=('y', 'mean'), n=('y', 'size'))
    imp = pd.Series(model().fit(prep(fin), fin.y).feature_importances_, index=FEATS).sort_values(ascending=False)
    print(dec.round(3).to_string()); print(now.p.describe().round(3).to_dict()); print(imp.head(8).to_dict())
    pickle.dump({'res': res, 'quintiles': dec, 'now': now, 'imp': imp, 'te': te, 'base_rate_final': fin.y.mean()}, open(f'{WORK}/churn.pkl', 'wb'))