"""Final next-best-SKU picks for every outlet that has bought in the window, scored on all data to 20 Oct 2026."""
import numpy as np, pandas as pd, pickle
from p03_reco import build, topk, B, K
from pcommon import *

T = pickle.load(open(f'{WORK}/reco_tuned.pkl', 'rb')); w = T['weights']; hit = T['hit_by_rank']
P = pd.read_parquet(f'{WORK}/products.parquet').set_index('sku')
F = build(DATA_END, with_label=False, scope='all')
r = lambda c: F.groupby('outlet_id')[c].rank(pct=True).to_numpy()
F['score'] = w['w_itemcf'] * r('itemcf') + w['w_van'] * r('vanshare') + w['w_usercf'] * r('usercf') + w['w_pop'] * r('pop') + w['w_aff'] * (r('prin_aff') + r('cat_aff')) + 1e-9 * F.itemcf
# only pick SKUs that sold somewhere in the last 4 weeks (still on the vans); in the 25 Aug back-test this filter changed no pick
last_sale = B.groupby('sku').date.max()
live = F.sku.isin(last_sale[last_sale > DATA_END - pd.Timedelta(days=28)].index)
print('candidate rows dropped by the 4-week availability filter:', int((~live).sum()))
picks = topk(F[live], F.score[live])
# typical spend in the first 8 weeks after an outlet starts buying a SKU (adoptions after 1 Apr by outlets already buying for 4+ weeks)
first_o = B.groupby('outlet_id').date.min()
fp = B.groupby(['outlet_id', 'sku']).date.min().rename('d0').reset_index()
fp = fp[(fp.d0 >= '2026-04-01') & (fp.d0 >= fp.outlet_id.map(first_o) + pd.Timedelta(days=28)) & (fp.d0 <= DATA_END - pd.Timedelta(days=56))]
x = B.merge(fp, on=['outlet_id', 'sku']); x = x[(x.date >= x.d0) & (x.date < x.d0 + pd.Timedelta(days=56))]
ad = x.groupby(['outlet_id', 'sku']).net_amt.sum().reset_index()
sku_med = ad.groupby('sku').net_amt.agg(['median', 'size'])
cat_med = ad.assign(c=ad.sku.map(P.l4_desc)).groupby('c').net_amt.median()
glob = ad.net_amt.median()
def typical(s):
    if s in sku_med.index and sku_med.loc[s, 'size'] >= 5: return sku_med.loc[s, 'median']
    c = P.l4_desc.get(s); return cat_med.get(c, glob)
picks['p_buy_8w'] = picks['rank'].map(hit).round(4)
picks['typical_8w_spend_sar'] = picks.sku.map(typical).round(2)
picks['expected_value_8w_sar'] = (picks.p_buy_8w * picks.typical_8w_spend_sar).round(2)
def reason(rw):
    if rw.vanshare >= 0.15: return f"{rw.vanshare:.0%} of outlets on van {rw.mainvan} buy it"
    if rw.usercf >= 0.15: return f"{rw.usercf:.0%} of outlets with a similar basket buy it"
    return f"popular: {rw['pop']:.0%} of active outlets buy it"
picks['why'] = picks.apply(reason, axis=1)
picks['sku_name'] = picks.sku.map(P.prd_desc_en); picks['principal'] = picks.sku.map(P.l1_desc)
out = picks[['outlet_id', 'rank', 'sku', 'sku_name', 'principal', 'p_buy_8w', 'typical_8w_spend_sar', 'expected_value_8w_sar', 'why', 'mainvan']].rename(columns={'mainvan': 'main_van'})
out['picks_made_on'] = DATA_END
out.to_parquet(f'{WORK}/reco_final.parquet', index=False)
print(out.shape, out.outlet_id.nunique(), out.head(6).to_string())