"""Tune blend weights of the next-best-SKU score on the validation cut-off (30 Jun), then score the 25 Aug test once."""
import itertools, numpy as np, pandas as pd, pickle
from p03_reco import build, topk, evaluate
from pcommon import *

val = build(pd.Timestamp('2026-06-30'))
te = pickle.load(open(f'{WORK}/reco_test.pkl', 'rb'))['te']
COMP = ['itemcf', 'vanshare', 'usercf', 'pop', 'prin_aff', 'cat_aff']

def ranks(F):
    return {c: F.groupby('outlet_id')[c].rank(pct=True).to_numpy() for c in COMP}

rv = ranks(val)
grid = []
for w in itertools.product([0, 0.2, 0.4, 0.6], repeat=4):              # itemcf, vanshare, usercf, pop
    if sum(w) == 0: continue
    for wa in (0, 0.1):                                                   # small principal / category affinity term
        sc = w[0] * rv['itemcf'] + w[1] * rv['vanshare'] + w[2] * rv['usercf'] + w[3] * rv['pop'] + wa * (rv['prin_aff'] + rv['cat_aff'])
        e = evaluate(val, topk(val, sc + 1e-9 * val.itemcf.to_numpy()))
        grid.append({'w_itemcf': w[0], 'w_van': w[1], 'w_usercf': w[2], 'w_pop': w[3], 'w_aff': wa, **e})
G = pd.DataFrame(grid).sort_values('precision_at_5', ascending=False)
print(G.head(8).round(4).to_string())
best = G.iloc[0]
rt = ranks(te)
sc = best.w_itemcf * rt['itemcf'] + best.w_van * rt['vanshare'] + best.w_usercf * rt['usercf'] + best.w_pop * rt['pop'] + best.w_aff * (rt['prin_aff'] + rt['cat_aff'])
te['score'] = sc + 1e-9 * te.itemcf.to_numpy()
picks = topk(te, te.score)
res = evaluate(te, picks)
print('TEST with tuned weights', {k: round(float(v), 4) for k, v in res.items()})
hit_by_rank = picks.groupby('rank').y.mean()
print(hit_by_rank.round(4).to_dict())
pickle.dump({'weights': best.to_dict(), 'grid': G, 'test_res': res, 'hit_by_rank': hit_by_rank, 'picks': picks}, open(f'{WORK}/reco_tuned.pkl', 'wb'))