"""Cluster adoption: will an outlet start buying from a sales team (cluster) it has never bought from?

Replicates the site's set-up on the five clusters sold under distributor 1000 (Personal Care sells under its own
codes, which the site merges by location and name; it is left out here). 30% of outlets held out.
Validation: train on pairs at 15 Apr (label 16 Apr - 30 Jun), validate on held-out outlets at 15 May (label 16 May - 31 Jul).
Test: train on 15 Apr + 15 May pairs (labels inside Mar-Jul), score held-out outlets at 31 Jul (label 1 Aug - 20 Oct), once.
"""
import json, pickle, itertools, warnings
import numpy as np, pandas as pd, lightgbm as lgb
warnings.filterwarnings('ignore')
from sklearn.metrics import roc_auc_score
from sklearn.neighbors import BallTree
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from catboost import CatBoostClassifier
from xgboost import XGBClassifier

ROOT = '/home/claude/binzagr'
s = pd.read_parquet(f'{ROOT}/pred/work/sales.parquet')
V = json.load(open(f'{ROOT}/boost/site_vans.json')); VMAP = dict(V['all'])
s['cluster'] = s.van.map(VMAP)
s = s[s.sellable & (s.net_amt > 0) & (s.cluster != 'Personal Care')]
C = pd.read_parquet('/mnt/user-data/outputs/Binzagr_Data/02_Cleaned/data/dim_customer.parquet')
C['outlet_id'] = C.dist_cd.astype(str) + '-' + C.cust_cd.astype(str); C = C.set_index('outlet_id')
CL = sorted(s.cluster.unique()); R_EARTH = 6371.0
rng = np.random.RandomState(20261023); outs_all = s.outlet_id.unique()
HOLD = set(rng.choice(outs_all, size=int(0.3 * len(outs_all)), replace=False))


def pairs(state, label_end):
    past = s[s.date <= state]; fut = s[(s.date > state) & (s.date <= label_end)]
    outs = past.outlet_id.unique()
    bought = set(zip(past.outlet_id, past.cluster)); adopt = set(zip(fut.outlet_id, fut.cluster))
    o = pd.DataFrame(index=pd.Index(outs, name='outlet_id'))
    o['spend'] = past.groupby('outlet_id').net_amt.sum(); o['spend_8w'] = past[past.date > state - pd.Timedelta(weeks=8)].groupby('outlet_id').net_amt.sum()
    o['skus'] = past.groupby('outlet_id').sku.nunique(); o['days'] = past.groupby('outlet_id').date.nunique()
    o['n_clusters'] = past.groupby('outlet_id').cluster.nunique()
    o['tenure'] = (state - past.groupby('outlet_id').date.min()).dt.days; o['recency'] = (state - past.groupby('outlet_id').date.max()).dt.days
    o = o.fillna(0)
    cc = C.reindex(o.index)
    o['channel'] = cc.cust_hier2.fillna('NA').astype(str).values; o['credit'] = (cc.cust_type == 'Q').astype(int).values
    o['lat'] = cc.latitude.values; o['lon'] = cc.longitude.values
    rows = []
    xy = np.radians(o[['lat', 'lon']].values)
    for c in CL:
        buyers = o.index[[(x, c) in bought for x in o.index]]
        bxy = np.radians(o.loc[buyers, ['lat', 'lon']].values)
        tree = BallTree(bxy, metric='haversine')
        d, _ = tree.query(xy, k=2)
        own = np.array([(x, c) in bought for x in o.index])
        nearest = np.where(own, d[:, 1], d[:, 0]) * R_EARTH                     # exclude the outlet itself if it buys
        n500 = tree.query_radius(xy, r=0.5 / R_EARTH, count_only=True) - own
        all_tree = BallTree(xy, metric='haversine'); tot500 = all_tree.query_radius(xy, r=0.5 / R_EARTH, count_only=True) - 1
        ch_up = pd.Series(own, index=o.index).groupby(o.channel).transform('mean').values
        df = o.copy(); df['cluster'] = c; df['dist_km'] = nearest; df['buyers_500m'] = n500
        df['share_500m'] = n500 / np.maximum(tot500, 1); df['channel_uptake'] = ch_up; df['cluster_uptake'] = own.mean()
        df['own'] = own; df['y'] = [int((x, c) in adopt) for x in o.index]
        rows.append(df[~df.own])
    P = pd.concat(rows).reset_index().drop(columns='own')
    return P


FEATS = ['spend', 'spend_8w', 'skus', 'days', 'n_clusters', 'tenure', 'recency', 'credit', 'dist_km', 'buyers_500m', 'share_500m',
         'channel_uptake', 'cluster_uptake', 'channel', 'cluster']


def Xcat(P):
    X = P[FEATS].copy()
    for c in ('channel', 'cluster'): X[c] = pd.Categorical(X[c], categories=sorted(set(C.cust_hier2.fillna('NA').astype(str)) | {'NA'}) if c == 'channel' else CL)
    return X


def Xnum(P):
    X = P[[f for f in FEATS if f not in ('channel', 'cluster')]].copy()
    for c in ('spend', 'spend_8w', 'skus', 'days', 'tenure', 'recency', 'dist_km', 'buyers_500m'): X[c] = np.log1p(X[c].clip(lower=0))
    X = pd.concat([X, pd.get_dummies(P.cluster, prefix='cl').astype(float)], axis=1)
    return X


def fit_predict(name, tr, te):
    if name == 'LightGBM':
        m = lgb.LGBMClassifier(n_estimators=400, 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)
        return m.fit(Xcat(tr), tr.y).predict_proba(Xcat(te))[:, 1]
    if name == 'Logistic regression':
        a, b = Xnum(tr), Xnum(te).reindex(columns=Xnum(tr).columns, fill_value=0)
        return make_pipeline(StandardScaler(), LogisticRegression(C=0.5, max_iter=3000)).fit(a, tr.y).predict_proba(b)[:, 1]
    if name == 'CatBoost':
        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(tr[FEATS], tr.y, cat_features=['channel', 'cluster']).predict_proba(te[FEATS])[:, 1]
    if name == 'XGBoost':
        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(Xcat(tr), tr.y).predict_proba(Xcat(te))[:, 1]
    raise KeyError(name)


BASE = ['LightGBM', 'Logistic regression', 'CatBoost', 'XGBoost']


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


def scores(tr, te):
    sc = {m: fit_predict(m, tr, te) for m in BASE}
    for k in (2, 3, 4):
        for combo in itertools.combinations(BASE, k): sc['Blend: ' + ' + '.join(combo)] = np.mean([rank(sc[m]) for m in combo], 0)
    sc['Channel popularity'] = te.channel_uptake.values; sc['What neighbours buy'] = te.share_500m.values
    return sc


def metrics(te, v):
    v = np.asarray(v, float); out = {'auc_pooled': roc_auc_score(te.y, v)}
    w = [roc_auc_score(te.y[te.cluster == c], v[(te.cluster == c).values]) for c in CL if te.y[te.cluster == c].nunique() == 2]
    out['auc_within_mean'] = float(np.mean(w))
    n = int(np.ceil(0.1 * len(te))); top = np.argsort(-v)[:n]; out['top10_capture'] = te.y.values[top].sum() / te.y.sum()
    return out


if __name__ == '__main__':
    A = pairs(pd.Timestamp('2026-04-15'), pd.Timestamp('2026-06-30'))
    Bv = pairs(pd.Timestamp('2026-05-15'), pd.Timestamp('2026-07-31'))
    T = pairs(pd.Timestamp('2026-07-31'), pd.Timestamp('2026-10-20'))
    tr_v, va = A[~A.outlet_id.isin(HOLD)], Bv[Bv.outlet_id.isin(HOLD)]
    tr_t, te = pd.concat([A[~A.outlet_id.isin(HOLD)], Bv[~Bv.outlet_id.isin(HOLD)]]), T[T.outlet_id.isin(HOLD)].reset_index(drop=True)
    va = va.reset_index(drop=True)
    print('validation pairs', len(va), 'adopted', round(va.y.mean(), 4), '| test pairs', len(te), 'adopted', int(te.y.sum()), round(te.y.mean(), 4), flush=True)
    print('test adoption rate by cluster', te.groupby('cluster').y.mean().round(4).to_dict())
    sv, st = scores(tr_v, va), scores(tr_t, te)
    rows = [{'model': k, **{f'val_{a}': b for a, b in metrics(va, sv[k]).items()}, **{f'test_{a}': b for a, b in metrics(te, st[k]).items()}} for k in st]
    R = pd.DataFrame(rows).sort_values('val_auc_within_mean', ascending=False)
    pick = R[~R.model.isin(['Channel popularity', 'What neighbours buy'])].iloc[0].model
    print(R.round(4).to_string(index=False)); print('picked on validation (within-cluster AUC):', pick)
    pickle.dump({'table': R, 'pick': pick, 'n_test': len(te), 'pos': int(te.y.sum())}, open(f'{ROOT}/boost/work/adoption_race.pkl', 'wb'))