"""Next-best SKU: replicate the site's test and compare recommenders.

Site protocol (Next-best SKU page): ten SKUs per outlet that it has not bought before the cut-off, judged on the
eight weeks after it; precision@10 = picks bought / (10 x outlets); 4,097 outlets = outlets that bought before the
cut-off and bought anything in the eight weeks after. Site: 12.0% vs 3.95% best-sellers vs 0.66% random.

Validation cut-off 30 Jun (judged 1 Jul - 25 Aug) for every choice; test cut-off 25 Aug (judged 26 Aug - 20 Oct), once.
"""
import sys, pickle, time, warnings
import numpy as np, pandas as pd, scipy.sparse as sp, lightgbm as lgb
warnings.filterwarnings('ignore')
sys.path.insert(0, '/home/claude/binzagr/pred')
from p03_reco import build, B

OUT = '/home/claude/binzagr/boost/work'
H = pd.Timedelta(days=56); K = 10
VAL, TEST, TRAIN0 = pd.Timestamp('2026-06-30'), pd.Timestamp('2026-08-25'), pd.Timestamp('2026-06-02')


def frame(cut):
    """Candidates (outlet x SKU new to it) for outlets that bought before the cut and buy again in the 8 weeks after."""
    F = build(cut, with_label=True, scope='all')
    fut = set(B[(B.date > cut) & (B.date <= cut + H)].outlet_id)
    F = F[F.outlet_id.isin(fut)].reset_index(drop=True)
    past = B[B.date <= cut]
    outs = np.array(sorted(F.outlet_id.unique())); skus = np.array(sorted(F.sku.unique()))
    oi = {o: i for i, o in enumerate(outs)}; si = {s_: j for j, s_ in enumerate(skus)}
    pp = past[past.outlet_id.isin(oi) & past.sku.isin(si)].drop_duplicates(['outlet_id', 'sku'])
    X = sp.csr_matrix((np.ones(len(pp)), (pp.outlet_id.map(oi).values, pp.sku.map(si).values)), shape=(len(outs), len(skus)))
    # recency-weighted copy (90-day half-life style decay), capped
    pw = past[past.outlet_id.isin(oi) & past.sku.isin(si)]
    w = np.exp(-(cut - pw.date).dt.days / 90.0)
    Xw = sp.csr_matrix((w.values, (pw.outlet_id.map(oi).values, pw.sku.map(si).values)), shape=X.shape); Xw.data = np.minimum(Xw.data, 3.0)
    F['oi'] = F.outlet_id.map(oi).values; F['si'] = F.sku.map(si).values; F['ocode'] = F['oi']
    return F, X, Xw


def ease(X, lam):
    G = (X.T @ X).toarray().astype(float); G[np.diag_indices_from(G)] += lam
    Pm = np.linalg.inv(G); Bm = -Pm / np.diag(Pm); np.fill_diagonal(Bm, 0)
    return Bm


def rp3(X, alpha, beta):
    Pui = sp.diags(1 / np.maximum(np.asarray(X.sum(1)).ravel(), 1)) @ X
    Piu = sp.diags(1 / np.maximum(np.asarray(X.sum(0)).ravel(), 1)) @ X.T
    Wm = (Piu @ Pui).toarray() ** alpha
    pop = np.maximum(np.asarray(X.sum(0)).ravel(), 1)
    Wm = Wm / (pop[None, :] ** beta); np.fill_diagonal(Wm, 0)
    return Wm


def score_matrix(X, Wm, F):
    S = np.asarray(X @ Wm)
    return S[F.oi.values, F.si.values]


def evaluate(F, score, k=K):
    """precision@10 / @5 and share of outlets with a hit."""
    oc = F['ocode'].values; y = F.y.values
    order = np.lexsort((-np.asarray(score, float), oc))
    y, oc = y[order], oc[order]
    start = np.r_[0, np.flatnonzero(np.diff(oc)) + 1]
    rank = np.arange(len(oc)) - np.repeat(start, np.diff(np.r_[start, len(oc)])) + 1
    n = len(start)
    hit_out = np.add.reduceat(y * (rank <= k), start) > 0
    return {'p@10': y[rank <= k].sum() / (k * n), 'p@5': y[rank <= 5].sum() / (5 * n), 'outlets_hit': hit_out.mean(), 'n': n}


def random_precision(F):
    return F.groupby('outlet_id').y.mean().mean()


def rank_pct(F, col):
    return F.groupby('outlet_id')[col].rank(pct=True).to_numpy()


if __name__ == '__main__':
    t0 = time.time()
    data = {c: frame(c) for c in (TRAIN0, VAL, TEST)}
    print('frames', {str(c.date()): (d[0].outlet_id.nunique(), len(d[0])) for c, d in data.items()}, f'{time.time() - t0:.0f}s', flush=True)
    T = pickle.load(open('/home/claude/binzagr/pred/work/reco_tuned.pkl', 'rb'))['weights']
    res = {}
    for c, (F, X, Xw) in data.items():
        F['site_blend'] = (T['w_itemcf'] * rank_pct(F, 'itemcf') + T['w_van'] * rank_pct(F, 'vanshare') + T['w_usercf'] * rank_pct(F, 'usercf')
                           + T['w_pop'] * rank_pct(F, 'pop') + T['w_aff'] * (rank_pct(F, 'prin_aff') + rank_pct(F, 'cat_aff')) + 1e-9 * F.itemcf)
    # ---- tune EASE and RP3beta on the validation cut-off
    Fv, Xv, Xwv = data[VAL]
    grid_e = {lam: evaluate(Fv, score_matrix(Xv, ease(Xv, lam), Fv))['p@10'] for lam in (10, 30, 100, 300, 1000, 3000)}
    lam = max(grid_e, key=grid_e.get); print('EASE lambda grid (val p@10)', {k_: round(v, 4) for k_, v in grid_e.items()}, '->', lam, flush=True)
    grid_r = {(a, b): evaluate(Fv, score_matrix(Xv, rp3(Xv, a, b), Fv))['p@10'] for a in (0.6, 0.8, 1.0, 1.2) for b in (0.0, 0.2, 0.4, 0.6)}
    ab = max(grid_r, key=grid_r.get); print('RP3beta grid best', ab, round(grid_r[ab], 4), flush=True)
    for c, (F, X, Xw) in data.items():
        F['ease'] = score_matrix(X, ease(X, lam), F)
        F['rp3'] = score_matrix(X, rp3(X, *ab), F)
    # ---- rank blends tuned on validation (EASE + RP3 + site blend + popularity)
    for c, (F, X, Xw) in data.items():
        for col in ('ease', 'rp3', 'site_blend', 'pop', 'vanshare'):
            F[f'r_{col}'] = rank_pct(F, col)
    best, grid_b = None, []
    for we in (0, 0.25, 0.5, 1):
        for wr in (0, 0.25, 0.5, 1):
            for ws in (0, 0.5, 1):
                for wv in (0, 0.25):
                    if we + wr + ws + wv == 0: continue
                    sc = we * Fv.r_ease + wr * Fv.r_rp3 + ws * Fv.r_site_blend + wv * Fv.r_vanshare + 1e-9 * Fv.ease
                    p = evaluate(Fv, sc)['p@10']; grid_b.append((p, we, wr, ws, wv))
    grid_b.sort(reverse=True); pb, we, wr, ws, wv = grid_b[0]
    print('blend val best p@10', round(pb, 4), 'weights ease/rp3/site/van', we, wr, ws, wv, flush=True)
    # ---- LightGBM ranker on every feature + the new scores (trained on the 2 Jun cut-off for validation; 2 Jun + 30 Jun for test)
    FEATS = ['itemcf', 'usercf', 'vanshare', 'pop', 'prin_aff', 'cat_aff', 'log_price', 'o_sales', 'o_skus', 'o_days', 'ease', 'rp3', 'site_blend',
             'r_ease', 'r_rp3', 'r_site_blend', 'r_pop']
    def fit_ranker(frames):
        D = pd.concat(frames, ignore_index=True)
        rs = np.random.RandomState(11); D = D[(D.y == 1) | (rs.rand(len(D)) < 0.15)].sort_values('outlet_id', kind='stable')
        grp = D.groupby('outlet_id', sort=False).size().values
        mdl = lgb.LGBMRanker(objective='lambdarank', n_estimators=300, 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, lambdarank_truncation_level=20)
        return mdl.fit(D[FEATS], D.y, group=grp)
    F0 = data[TRAIN0][0].assign(outlet_id=lambda d: 'a' + d.outlet_id)
    rk_val = fit_ranker([F0]); Fv['ranker'] = rk_val.predict(Fv[FEATS])
    F1 = Fv.assign(outlet_id=lambda d: 'b' + d.outlet_id)
    rk_test = fit_ranker([F0, F1]); Ft = data[TEST][0]; Ft['ranker'] = rk_test.predict(Ft[FEATS])
    # ---- score every strategy on validation and test
    strategies = {
        'Best-sellers (popularity)': lambda F: F['pop'] + 1e-9 * F.itemcf,
        'Site-style blend (similar baskets + van + popularity)': lambda F: F.site_blend,
        'EASE (item model, closed form)': lambda F: F.ease,
        'RP3beta (graph walk)': lambda F: F.rp3,
        'Blend: EASE + RP3beta + site blend + van': lambda F: we * F.r_ease + wr * F.r_rp3 + ws * F.r_site_blend + wv * F.r_vanshare + 1e-9 * F.ease,
        'LightGBM ranker on all scores': lambda F: F.ranker,
    }
    rows = []
    for name, fn in strategies.items():
        if name.startswith('LightGBM') and 'ranker' not in Fv: continue
        v = evaluate(Fv, fn(Fv)); t = evaluate(Ft, fn(Ft))
        rows.append({'strategy': name, 'val_p10': v['p@10'], 'test_p10': t['p@10'], 'test_p5': t['p@5'], 'test_outlets_hit': t['outlets_hit'], 'test_outlets': t['n']})
    rows.append({'strategy': 'Random SKUs', 'val_p10': random_precision(Fv), 'test_p10': random_precision(Ft), 'test_outlets': Ft.outlet_id.nunique()})
    R = pd.DataFrame(rows)
    print(R.round(4).to_string(index=False))
    pickle.dump({'table': R, 'lam': lam, 'ab': ab, 'blend_w': (we, wr, ws, wv), 'grid_e': grid_e, 'grid_r': grid_r}, open(f'{OUT}/reco_race.pkl', 'wb'))
    print(f'done {time.time() - t0:.0f}s')