"""Next-best SKU per outlet: a LightGBM ranker over collaborative-filtering and van / principal affinity features.

Backtest exactly as the wireframe states: picks made on 25 Aug 2026 from data up to that day, judged on what each
outlet bought for the first time in the 8 unseen weeks after it (26 Aug - 20 Oct), against "push best-sellers" and random.
Training labels come from earlier cut-offs (2 Jun and 30 Jun, each followed by 8 weeks).
"""
import numpy as np, pandas as pd, scipy.sparse as sp, lightgbm as lgb, pickle
from sklearn.metrics import roc_auc_score
from pcommon import *

S = pd.read_parquet(f'{WORK}/sales.parquet'); P = pd.read_parquet(f'{WORK}/products.parquet').set_index('sku')
B = S[S.sellable & (S.net_amt > 0)].copy()
B['principal'] = B.sku.map(P.l1_code).fillna('NA'); B['cat4'] = B.sku.map(P.l4_desc).fillna('NA')
H = pd.Timedelta(days=56)
K = 5


def build(cut, with_label=True, scope='recent'):
    """Candidate (outlet, SKU) pairs at cut-off `cut` with features; label = first purchase in the next 8 weeks."""
    past = B[B.date <= cut]
    recent = past[past.date > cut - H]
    outlets = recent.outlet_id.unique() if scope == 'recent' else past.outlet_id.unique()   # 'recent' = active in the 8 weeks before the cut-off
    assort = recent.sku.unique()                                           # SKUs on the vans in the last 8 weeks
    o_ix = {o: i for i, o in enumerate(outlets)}; s_ix = {s: i for i, s in enumerate(assort)}
    pp = past[past.outlet_id.isin(o_ix) & past.sku.isin(s_ix)]
    # outlet x SKU purchase matrix (all history up to the cut-off), recency-weighted
    w = np.exp(-(cut - pp.date).dt.days / 90.0)
    M = sp.csr_matrix((w.values, (pp.outlet_id.map(o_ix).values, pp.sku.map(s_ix).values)), shape=(len(outlets), len(assort)))
    M.data = np.minimum(M.data, 3.0)
    Bin = (M > 0).astype(float)
    # item-item cosine similarity (co-purchase across outlets)
    norms = np.sqrt(np.asarray(Bin.multiply(Bin).sum(axis=0))).ravel() + 1e-9
    Xn = Bin.multiply(1 / norms).tocsc()
    simI = (Xn.T @ Xn).toarray(); np.fill_diagonal(simI, 0)
    itemcf = np.asarray(M @ simI)
    itemcf = itemcf / (np.asarray(M.sum(axis=1)).ravel()[:, None] + 1e-9)
    # user-user: outlets with similar baskets (cosine), top-30 neighbours
    rn = np.sqrt(np.asarray(Bin.multiply(Bin).sum(axis=1))).ravel() + 1e-9
    Un = sp.diags(1 / rn) @ Bin
    simU = (Un @ Un.T).toarray(); np.fill_diagonal(simU, 0)
    if simU.shape[0] > 30:
        thr = -np.partition(-simU, 30, axis=1)[:, 30:31]
        simU = np.where(simU >= thr, simU, 0)
    usercf = (simU @ Bin.toarray()) / (simU.sum(axis=1, keepdims=True) + 1e-9)
    # van affinity: share of the outlet's main van's customers that bought the SKU in the last 8 weeks
    mainvan = recent.groupby('outlet_id').van.agg(lambda x: x.value_counts().index[0])
    lastvan = past.sort_values('date').groupby('outlet_id').van.last()
    mainvan = mainvan.reindex(outlets).fillna(lastvan.reindex(outlets))
    vo = recent.groupby(['van', 'sku']).outlet_id.nunique().unstack(fill_value=0)
    vn = recent.groupby('van').outlet_id.nunique()
    vanshare = (vo.div(vn, axis=0)).reindex(index=mainvan.reindex(outlets).values, columns=assort).fillna(0).to_numpy()
    # global popularity (outlets buying in last 8 weeks)
    pop = recent.groupby('sku').outlet_id.nunique().reindex(assort).fillna(0).to_numpy() / len(outlets)
    # outlet affinity to the SKU's principal and category
    def aff(col):
        t = past.groupby(['outlet_id', col]).net_amt.sum().unstack(fill_value=0)
        t = t.div(t.sum(axis=1), axis=0).reindex(outlets).fillna(0)
        keyv = pd.Series(assort).map(B.drop_duplicates('sku').set_index('sku')[col]).fillna('NA').values
        return t.reindex(columns=pd.Index(keyv)).fillna(0).to_numpy()
    prin_aff, cat_aff = aff('principal'), aff('cat4')
    price = (recent.groupby('sku').net_amt.sum() / recent.groupby('sku').qty_cases.sum()).reindex(assort).fillna(0).to_numpy()
    o_sales = recent.groupby('outlet_id').net_amt.sum().reindex(outlets).fillna(0).to_numpy()
    o_skus = past.groupby('outlet_id').sku.nunique().reindex(outlets).to_numpy()
    o_days = recent.groupby('outlet_id').date.nunique().reindex(outlets).fillna(0).to_numpy()
    bought = Bin.toarray() > 0
    oi, si = np.where(~bought)                                               # candidates: SKUs new to the outlet
    F = pd.DataFrame({'outlet_id': outlets[oi], 'sku': assort[si], 'itemcf': itemcf[oi, si], 'usercf': usercf[oi, si], 'vanshare': vanshare[oi, si],
                      'pop': pop[si], 'prin_aff': prin_aff[oi, si], 'cat_aff': cat_aff[oi, si], 'log_price': np.log1p(price[si]),
                      'o_sales': np.log1p(o_sales[oi]), 'o_skus': o_skus[oi], 'o_days': o_days[oi]})
    F['mainvan'] = F.outlet_id.map(mainvan)
    if with_label:
        fut = B[(B.date > cut) & (B.date <= cut + H)]
        got = set(zip(fut.outlet_id, fut.sku))
        F['y'] = [int((o, s) in got) for o, s in zip(F.outlet_id, F.sku)]
    return F


FEATS = ['itemcf', 'usercf', 'vanshare', 'pop', 'prin_aff', 'cat_aff', 'log_price', 'o_sales', 'o_skus', 'o_days']


def topk(F, score, k=K):
    F = F.assign(score=score).sort_values(['outlet_id', 'score'], ascending=[True, False])
    F['rank'] = F.groupby('outlet_id').cumcount() + 1
    return F[F['rank'] <= k]


def evaluate(F, picks):
    hit = picks.y.mean()
    per = picks.groupby('outlet_id').y.max().mean()
    rnd = (F.groupby('outlet_id').y.mean()).mean()                       # expected precision of 5 random new SKUs
    return {'precision_at_5': hit, 'outlets_with_hit': per, 'random_precision': rnd}


if __name__ == '__main__':
    tr = pd.concat([build(pd.Timestamp('2026-06-02')), build(pd.Timestamp('2026-06-30'))], ignore_index=True)
    te = build(RECO_CUTOFF)
    print('train rows', len(tr), 'positives', tr.y.mean().round(4), '| test rows', len(te), 'positives', te.y.mean().round(4))
    rk = lgb.LGBMClassifier(n_estimators=400, learning_rate=0.05, num_leaves=63, min_child_samples=100, subsample=0.8, subsample_freq=1,
                            colsample_bytree=0.8, verbose=-1, n_jobs=2)
    rk.fit(tr[FEATS], tr.y)
    te['p'] = rk.predict_proba(te[FEATS])[:, 1]
    res = {}
    res['next_best_sku_model'] = evaluate(te, topk(te, te.p))
    res['item_cf_only'] = evaluate(te, topk(te, te.itemcf))
    res['push_best_sellers'] = evaluate(te, topk(te, te['pop'] + 1e-6 * te.itemcf))
    res['van_best_sellers'] = evaluate(te, topk(te, te.vanshare + 1e-6 * te['pop']))
    res['auc_model'] = roc_auc_score(te.y, te.p)
    for k, v in res.items(): print(k, v)
    pickle.dump({'model': rk, 'res': res, 'te': te, 'tr_cols': FEATS}, open(f'{WORK}/reco_backtest.pkl', 'wb'))