"""Van rep day: suggested van load (van x SKU forecast), next-week visit plan, van productivity."""
import numpy as np, pandas as pd, lightgbm as lgb, pickle
from sklearn.metrics import roc_auc_score
from pcommon import *
import fc_engine as E

S = pd.read_parquet(f'{WORK}/sales.parquet'); R = pd.read_parquet(f'{WORK}/returns.parquet'); P = pd.read_parquet(f'{WORK}/products.parquet').set_index('sku')

# ------------------------------------------------------------------ 1. suggested load: van x SKU weekly cases
S2 = S[(S.week_end >= FIRST_FULL_WEEK_END)]
recent_pairs = S2[S2.week_end > pd.Timestamp('2026-06-23')].groupby(['van', 'sku']).size().index
S2 = S2.set_index(['van', 'sku']).loc[lambda d: d.index.isin(recent_pairs)].reset_index()
S2['uid'] = S2.van + '|' + S2.sku
df = S2.groupby(['uid', 'week_end']).qty_cases.sum().reset_index().rename(columns={'uid': 'unique_id', 'week_end': 'ds', 'qty_cases': 'y'})
ids_all = df.unique_id.unique()
grid = pd.MultiIndex.from_product([ids_all, E.WEEKS], names=['unique_id', 'ds']).to_frame(index=False)
df = grid.merge(df, how='left').fillna({'y': 0.0})
ids, Y = E._panel(df)
cats = np.array([P.l1_code.get(u.split('|')[1], 'NA') for u in ids])
print('van x sku series', len(ids))
# hold-out: train to 22 Sep, rolling one week ahead over the 4 hold-out weeks
mods, lev = E.lgb_fit(Y, cats, E.widx(TRAIN_END), horizons=[1], alphas=(0.1, 0.5, 0.9), rounds=300)
rows = []
for tw in E.TEST_T:
    t = E.widx(tw) - 1
    p, sc = E.lgb_predict(mods, lev, Y, cats, t, 1)
    rows.append(pd.DataFrame({'unique_id': ids, 'ds': tw, 'y': Y[:, t + 1], 'naive': Y[:, t], 'p10': p[0.1], 'p50': p[0.5], 'p90': p[0.9]}))
ho = pd.concat(rows, ignore_index=True)
ho['van'] = ho.unique_id.str.split('|').str[0]
act = ho[ho.groupby('unique_id').y.transform('sum') + ho.groupby('unique_id').naive.transform('sum') > 0]
vt = ho.groupby(['van', 'ds'])[['y', 'naive', 'p50', 'p90']].sum()
load_metrics = {'series': len(ids), 'holdout_wape_p50_van_sku': wape(act.y, act.p50), 'holdout_wape_naive_van_sku': wape(act.y, act.naive),
                'holdout_share_weeks_covered_by_p90_load': float((act.y <= act.p90 + 1e-9).mean()),
                'holdout_share_cases_covered_by_p90_load': float(np.minimum(act.y, act.p90).sum() / act.y.sum()),
                'holdout_p90_load_vs_sales': float(act.p90.sum() / act.y.sum()),
                'holdout_naive_load_vs_sales': float(act.naive.sum() / act.y.sum()),
                'holdout_share_cases_covered_by_naive_load': float(np.minimum(act.y, act.naive).sum() / act.y.sum()),
                'holdout_wape_p50_van_total_cases': wape(vt.y, vt.p50), 'holdout_wape_naive_van_total_cases': wape(vt.y, vt.naive)}
print({k: round(v, 3) for k, v in load_metrics.items()})
# next week: train on everything to 20 Oct
mods, lev = E.lgb_fit(Y, cats, E.widx(DATA_END), horizons=[1], alphas=(0.1, 0.5, 0.9), rounds=300)
p, sc = E.lgb_predict(mods, lev, Y, cats, E.widx(DATA_END), 1)
nxt = pd.DataFrame({'unique_id': ids, 'p10': p[0.1], 'p50': p[0.5], 'p90': p[0.9], 'last_week_sold_cases': Y[:, -1], 'avg_weekly_cases_8w': Y[:, -8:].mean(axis=1)})
nxt[['van', 'sku']] = nxt.unique_id.str.split('|', expand=True)
ret4 = R[R.date > DATA_END - pd.Timedelta(days=28)].groupby(['van', 'sku']).qty_cases.sum()
nxt['returns_last_4w_cases'] = [ret4.get((v, s), 0.0) for v, s in zip(nxt.van, nxt.sku)]
active_vans = S[S.date > DATA_END - pd.Timedelta(days=14)].van.unique()
nxt = nxt[nxt.van.isin(active_vans) & ((nxt.p90 > 0.05) | (nxt.last_week_sold_cases > 0))]
pickle.dump({'next': nxt, 'metrics': load_metrics, 'holdout': ho}, open(f'{WORK}/van_load.pkl', 'wb'))
print('next-week load rows', len(nxt), 'vans', nxt.van.nunique())

# ------------------------------------------------------------------ 2. visit plan: will van v visit outlet o on day d?
V = S.groupby(['van', 'outlet_id', 'date']).size().reset_index()[['van', 'outlet_id', 'date']]
WORKDAYS = lambda wk_end: [d for d in pd.date_range(wk_end - pd.Timedelta(days=6), wk_end) if d.weekday() != 4]   # Friday off


def visit_frame(week_end_):
    start = week_end_ - pd.Timedelta(days=6); asof = start - pd.Timedelta(days=1)
    h = V[V.date <= asof]
    pairs = h[h.date > asof - pd.Timedelta(weeks=8)][['van', 'outlet_id']].drop_duplicates()
    g = h.merge(pairs, on=['van', 'outlet_id'])
    f = pairs.copy()
    for wk in (4, 8, 13):
        x = g[g.date > asof - pd.Timedelta(weeks=wk)].groupby(['van', 'outlet_id']).date.nunique().rename(f'visits_{wk}w')
        f = f.merge(x, on=['van', 'outlet_id'], how='left')
    last = g.groupby(['van', 'outlet_id']).date.max().rename('last_visit'); f = f.merge(last, on=['van', 'outlet_id'], how='left')
    gaps = g[g.date > asof - pd.Timedelta(weeks=13)].groupby(['van', 'outlet_id']).date.apply(lambda s: np.median(np.diff(np.sort(s.unique()).astype('datetime64[D]').astype(int))) if s.nunique() > 1 else np.nan).rename('median_gap')
    f = f.merge(gaps, on=['van', 'outlet_id'], how='left')
    ov = V[(V.date <= asof) & (V.date > asof - pd.Timedelta(weeks=8))]
    f = f.merge(ov.groupby('outlet_id').date.nunique().rename('outlet_visit_days_8w'), on='outlet_id', how='left')
    vdays = ov.groupby('van').date.nunique().rename('van_days_8w'); f = f.merge(vdays, on='van', how='left')
    vact2 = V[(V.date <= asof) & (V.date > asof - pd.Timedelta(days=14))].groupby('van').date.nunique().rename('van_days_2w'); f = f.merge(vact2, on='van', how='left')
    rows = []
    g8 = g[g.date > asof - pd.Timedelta(weeks=8)].assign(dow=lambda d: d.date.dt.weekday)
    dowc = g8.groupby(['van', 'outlet_id', 'dow']).date.nunique().rename('same_dow_visits_8w')
    vdow = ov.assign(dow=ov.date.dt.weekday).groupby(['van', 'dow']).date.nunique().rename('van_days_on_dow_8w')
    for d in WORKDAYS(week_end_):
        x = f.copy(); x['date'] = d; x['dow'] = d.weekday()
        x = x.merge(dowc.reset_index(), on=['van', 'outlet_id', 'dow'], how='left').merge(vdow.reset_index(), on=['van', 'dow'], how='left')
        x['days_since_last'] = (d - x.last_visit).dt.days
        rows.append(x)
    F = pd.concat(rows, ignore_index=True).fillna({'same_dow_visits_8w': 0, 'van_days_on_dow_8w': 0, 'van_days_2w': 0})
    F['since_vs_gap'] = F.days_since_last / F.median_gap.fillna(28)
    F['same_dow_share'] = F.same_dow_visits_8w / F.visits_8w.clip(lower=1)
    F['van_share_of_outlet'] = F.visits_8w / F.outlet_visit_days_8w.clip(lower=1)
    hit = set(zip(V.van, V.outlet_id, V.date))
    F['y'] = [int((a, b, c) in hit) for a, b, c in zip(F.van, F.outlet_id, F.date)]
    return F


VF = ['dow', 'visits_4w', 'visits_8w', 'visits_13w', 'median_gap', 'outlet_visit_days_8w', 'van_days_8w', 'van_days_2w', 'same_dow_visits_8w',
      'van_days_on_dow_8w', 'days_since_last', 'since_vs_gap', 'same_dow_share', 'van_share_of_outlet']
weeks_tr = pd.date_range('2026-07-14', '2026-10-13', freq='W-TUE')
TR = pd.concat([visit_frame(w) for w in weeks_tr], ignore_index=True)
TE = visit_frame(pd.Timestamp('2026-10-20'))
vm = lgb.LGBMClassifier(n_estimators=300, learning_rate=0.05, num_leaves=31, min_child_samples=50, subsample=0.8, subsample_freq=1, colsample_bytree=0.8, verbose=-1, n_jobs=2)
vm.fit(TR[VF], TR.y)
TE['p'] = vm.predict_proba(TE[VF])[:, 1]
rule = TE.same_dow_visits_8w / 8                                           # rule: share of the last 8 same weekdays the van visited
visit_metrics = {'holdout_week': '2026-10-14 to 2026-10-20', 'candidate_van_outlet_days': len(TE), 'actual_visits': int(TE.y.sum()),
                 'auc_model': roc_auc_score(TE.y, TE.p), 'auc_same_weekday_rule': roc_auc_score(TE.y, rule)}
for thr in (0.3, 0.5):
    pr = TE.p >= thr
    visit_metrics[f'precision_at_{thr}'] = float(TE.y[pr].mean()); visit_metrics[f'recall_at_{thr}'] = float(TE.y[pr].sum() / TE.y.sum())
visit_metrics['predicted_visits_sum_p'] = float(TE.p.sum())
print({k: (round(v, 3) if isinstance(v, float) else v) for k, v in visit_metrics.items()})
TR2 = pd.concat([TR, TE.drop(columns='p')], ignore_index=True)
vm2 = lgb.LGBMClassifier(n_estimators=300, learning_rate=0.05, num_leaves=31, min_child_samples=50, subsample=0.8, subsample_freq=1, colsample_bytree=0.8, verbose=-1, n_jobs=2).fit(TR2[VF], TR2.y)
NX = visit_frame(pd.Timestamp('2026-10-27')); NX['p'] = vm2.predict_proba(NX[VF])[:, 1]

# ------------------------------------------------------------------ 3. van productivity (actual)
vd = S.groupby(['van', 'date']).agg(outlets=('outlet_id', 'nunique'), sales=('net_amt', 'sum')).reset_index()
last8 = vd[vd.date > DATA_END - pd.Timedelta(weeks=8)]
o_vans = S.groupby('outlet_id').van.nunique()
prod = vd.groupby('van').agg(active_days=('date', 'nunique'), first_day=('date', 'min'), last_day=('date', 'max'), median_outlets_per_day=('outlets', 'median'),
                             sales_sar=('sales', 'sum'), median_sales_per_day_sar=('sales', 'median'))
prod = prod.join(last8.groupby('van').agg(active_days_8w=('date', 'nunique'), median_outlets_per_day_8w=('outlets', 'median'), sales_8w_sar=('sales', 'sum')))
srv = S[S.date > DATA_END - pd.Timedelta(weeks=8)].groupby('van').outlet_id.nunique().rename('outlets_served_8w'); prod = prod.join(srv)
shared = S[S.date > DATA_END - pd.Timedelta(weeks=8)][['van', 'outlet_id']].drop_duplicates()
shared['multi'] = shared.outlet_id.map(S[S.date > DATA_END - pd.Timedelta(weeks=8)].groupby('outlet_id').van.nunique()) > 1
prod = prod.join(shared.groupby('van').multi.mean().rename('share_outlets_also_served_by_other_vans_8w'))
prod['active_now'] = prod.index.isin(active_vans)
prod_stats = {'median_outlets_per_van_day': float(vd.outlets.median()), 'best_quartile_outlets_per_van_day': float(prod.median_outlets_per_day.quantile(0.75)),
              'median_outlets_per_van_day_8w': float(last8.outlets.median()), 'best_quartile_outlets_per_van_day_8w': float(prod.median_outlets_per_day_8w.quantile(0.75)),
              'share_outlets_served_by_2plus_vans': float((o_vans >= 2).mean()), 'vans_active_last_14_days': int(len(active_vans)), 'vans_in_window': int(S.van.nunique())}
print({k: round(v, 3) if isinstance(v, float) else v for k, v in prod_stats.items()})
pickle.dump({'visit_next': NX, 'visit_metrics': visit_metrics, 'visit_holdout': TE, 'prod': prod, 'prod_stats': prod_stats}, open(f'{WORK}/van_day.pkl', 'wb'))