"""Next-best SKU, final step: fair pick between the EASE blend and the LightGBM ranker, bootstrap of the test gain,
and the improved picks for every outlet as of 20 Oct 2026.

Fair pick: blend weights tuned on the 2 Jun cut-off and scored at 30 Jun (the ranker was trained on 2 Jun and scored
at 30 Jun in reco_race.py, val p@10 0.1547), so both validation scores are out of sample.
"""
import sys, pickle, time, warnings
import numpy as np, pandas as pd, lightgbm as lgb
warnings.filterwarnings('ignore')
sys.path.insert(0, '/home/claude/binzagr/boost'); sys.path.insert(0, '/home/claude/binzagr/pred')
import reco_race as RR
from p03_reco import build, B, P
OUT = '/home/claude/binzagr/boost/work'
DATA_END = pd.Timestamp('2026-10-20')
RACE = pickle.load(open(f'{OUT}/reco_race.pkl', 'rb')); LAM, AB = RACE['lam'], RACE['ab']
T = pickle.load(open('/home/claude/binzagr/pred/work/reco_tuned.pkl', 'rb'))['weights']
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 enrich(F, X):
    F['site_blend'] = (T['w_itemcf'] * RR.rank_pct(F, 'itemcf') + T['w_van'] * RR.rank_pct(F, 'vanshare') + T['w_usercf'] * RR.rank_pct(F, 'usercf')
                       + T['w_pop'] * RR.rank_pct(F, 'pop') + T['w_aff'] * (RR.rank_pct(F, 'prin_aff') + RR.rank_pct(F, 'cat_aff')) + 1e-9 * F.itemcf)
    F['ease'] = RR.score_matrix(X, RR.ease(X, LAM), F); F['rp3'] = RR.score_matrix(X, RR.rp3(X, *AB), F)
    for col in ('ease', 'rp3', 'site_blend', 'pop', 'vanshare'): F[f'r_{col}'] = RR.rank_pct(F, col)
    return F


def blend_score(F, w):
    we, wr, ws, wv = w
    return we * F.r_ease + wr * F.r_rp3 + ws * F.r_site_blend + wv * F.r_vanshare + 1e-9 * F.ease


def tune(F):
    best = 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
                    p = RR.evaluate(F, blend_score(F, (we, wr, ws, wv)))['p@10']
                    if best is None or p > best[0]: best = (p, (we, wr, ws, wv))
    return best


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
    m = 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 m.fit(D[FEATS], D.y, group=grp)


def per_outlet_hits(F, score, k=10):
    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
    return np.add.reduceat(y * (rank <= k), start)


if __name__ == '__main__':
    t0 = time.time()
    fr = {c: RR.frame(c) for c in (RR.TRAIN0, RR.VAL, RR.TEST)}
    for c, (F, X, Xw) in fr.items(): enrich(F, X)
    F0, Fv, Ft = fr[RR.TRAIN0][0], fr[RR.VAL][0], fr[RR.TEST][0]
    # fair (out-of-sample) validation score of the blend: weights tuned at 2 Jun, scored at 30 Jun
    p0, w0 = tune(F0); fair_blend_val = RR.evaluate(Fv, blend_score(Fv, w0))['p@10']
    pv, wv = tune(Fv)                                    # weights used for the test (tuned on 30 Jun, as in the race)
    ranker_val = 0.1547                                  # reco_race.py: trained 2 Jun, scored 30 Jun
    pick = 'ranker' if ranker_val > fair_blend_val else 'blend'
    print(f'fair validation p@10: blend {fair_blend_val:.4f} (weights {w0}) vs ranker {ranker_val:.4f} -> pick {pick}; test weights {wv}', flush=True)
    # test scores and paired bootstrap over outlets
    rk = fit_ranker([F0.assign(outlet_id=lambda d: 'a' + d.outlet_id), Fv.assign(outlet_id=lambda d: 'b' + d.outlet_id)])
    scores = {'site_style': Ft.site_blend.values, 'ease': Ft.ease.values, 'blend': blend_score(Ft, wv).values, 'ranker': rk.predict(Ft[FEATS])}
    hits = {k: per_outlet_hits(Ft, v) for k, v in scores.items()}
    n = len(hits['site_style']); rs = np.random.RandomState(5); boot = {k: [] for k in ('ease', 'blend', 'ranker')}
    for _ in range(2000):
        i = rs.randint(0, n, n)
        for k in boot: boot[k].append((hits[k][i].sum() - hits['site_style'][i].sum()) / (10 * n))
    res = {k: {'test_p10': hits[k].sum() / (10 * n), 'gain_vs_site_style_p10': (hits[k].sum() - hits['site_style'].sum()) / (10 * n),
               'gain_90pct_interval': np.percentile(boot[k], [5, 95]).round(4).tolist() if k in boot else None,
               'outlets_hit': float((hits[k] > 0).mean())} for k in hits}
    for k, v in res.items(): print(k, {a: (round(b, 4) if isinstance(b, float) else b) for a, b in v.items()})
    # improved picks for every outlet that bought in the window, as of 20 Oct 2026, with the picked model
    Fn = build(DATA_END, with_label=False, scope='all'); past = B[B.date <= DATA_END]
    outs = np.array(sorted(Fn.outlet_id.unique())); skus = np.array(sorted(Fn.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'])
    import scipy.sparse as sp
    Xn = sp.csr_matrix((np.ones(len(pp)), (pp.outlet_id.map(oi).values, pp.sku.map(si).values)), shape=(len(outs), len(skus)))
    Fn['oi'] = Fn.outlet_id.map(oi).values; Fn['si'] = Fn.sku.map(si).values; Fn['ocode'] = Fn['oi']; Fn['y'] = 0
    enrich(Fn, Xn)
    last_sale = B.groupby('sku').date.max(); Fn = Fn[Fn.sku.isin(last_sale[last_sale > DATA_END - pd.Timedelta(days=28)].index)]
    if pick == 'ranker':
        rk_all = fit_ranker([F0.assign(outlet_id=lambda d: 'a' + d.outlet_id), Fv.assign(outlet_id=lambda d: 'b' + d.outlet_id),
                             Ft.assign(outlet_id=lambda d: 'c' + d.outlet_id)])
        Fn['score'] = rk_all.predict(Fn[FEATS])
    else:
        Fn['score'] = blend_score(Fn, wv).values
    Fn = Fn.sort_values(['outlet_id', 'score'], ascending=[True, False]); Fn['rank'] = Fn.groupby('outlet_id').cumcount() + 1
    top = Fn[Fn['rank'] <= 10][['outlet_id', 'rank', 'sku', 'score', 'mainvan']].copy()
    top['sku_name'] = top.sku.map(P.prd_desc_en); top['principal'] = top.sku.map(P.l1_desc)
    top['model'] = 'LightGBM ranker over EASE, RP3beta and the current scores' if pick == 'ranker' else f'Blend of EASE, RP3beta and the current scores {wv}'
    top.to_parquet(f'{OUT}/picks_v2.parquet', index=False)
    pickle.dump({'res': res, 'pick': pick, 'fair_blend_val': fair_blend_val, 'ranker_val': ranker_val, 'w_fair': w0, 'w_test': wv},
                open(f'{OUT}/reco_final_v2.pkl', 'wb'))
    print('picks', top.outlet_id.nunique(), 'outlets', len(top), 'rows', f'{time.time() - t0:.0f}s')