"""Warehouse labour screen: cost per activity vs best quartile, badge-sharing corrections, heads needed vs on shift,
role x activity hours, and a daily workload forecast for October 2026 (WMS data ends 30 Sep)."""
import sys, glob, numpy as np, pandas as pd, pickle
from pcommon import *
sys.path.insert(0, os.path.join(ROOT, '03_Simulated', 'scripts'))
from sim_labour import cycle_seconds, size_band, SYSTEM

cols = ['tran_type', 'tran_desc', 'tran_category', 'wh_id', 'start_ts', 'end_ts', 'dur_s', 'employee_id_norm', 'idle_secs_before',
        'person_session_id', 'person_id', 'wms_role', 'qty_cases', 'qty_pallets', 'is_billing_event', 'event_date']
T = pd.concat([pd.read_parquet(p, columns=cols) for p in sorted(glob.glob(f'{CLEAN}/jeddah_wms_transactions/*.parquet'))], ignore_index=True)
T = T[~T.employee_id_norm.isin(SYSTEM)].reset_index(drop=True)
T['tt'] = T.tran_type.astype(str)
T['cyc'] = cycle_seconds(T, 1800).astype('float64')
T['band'] = size_band(T.qty_cases)
st = pd.read_parquet(f'{SIM}/sim_standard_times.parquet'); st = st[st.cap_minutes == 30]
T['std_s'] = (T.tran_desc.astype(str) + '|' + T.band).map(st.assign(k=st.tran_desc + '|' + st.size_band).set_index('k').std_s_median_p50).fillna(0).astype(float)
T['wms_role'] = T.wms_role.astype(str).replace({'None': 'Not in user list', 'nan': 'Not in user list', '<NA>': 'Not in user list'})
for c in ('tran_desc', 'tran_category', 'wh_id'):
    T[c] = T[c].astype(str)
print('scans', len(T))

# ------------------------------------------------------------------ 1. cost per deck activity
def deck(r_tt, cat, qc, qp):
    out = np.full(len(r_tt), 'Other', dtype=object)
    out[np.isin(cat, ['Inbound'])] = 'Inbound receipt'
    out[np.isin(cat, ['Internal movement', 'Replenishment'])] = 'Pallet movement'
    out[np.isin(r_tt, ['203', '204', '205', '206'])] = 'Repalletisation'
    pick = np.isin(r_tt, ['301', '302'])
    out[pick] = 'Case picking'
    out[pick & (qc < 1)] = 'Piece picking'
    out[pick & (qp >= 1)] = 'Pallet picking'
    out[np.isin(cat, ['Loading / shipping']) & ~np.isin(r_tt, ['391', '392'])] = 'Truck loading'
    return out
T['activity'] = deck(T.tt.to_numpy(), T.tran_category.astype(str).to_numpy(), T.qty_cases.fillna(1).to_numpy(), T.qty_pallets.fillna(0).to_numpy())
PRIMARY = {'Inbound receipt': ['153', '183'], 'Pallet movement': ['201', '211', '255', '305'], 'Repalletisation': ['203', '205'],
           'Case picking': ['301'], 'Piece picking': ['301'], 'Pallet picking': ['301'], 'Truck loading': ['321']}
T['is_op'] = [tt in PRIMARY.get(a, []) for tt, a in zip(T.tt, T.activity)]
vol = pd.read_parquet(f'{SIM}/sim_activity_volumes.parquet')
vol['tt'] = vol.tran_type.astype(str)
rate = vol.groupby('tt').posted_labour_cost_sar.sum() / (vol.groupby('tt').cycle_hours.sum() * 3600)      # SAR per cycle-second, by type
hrs_rate = vol.groupby('tt').posted_paid_hours.sum() / (vol.groupby('tt').cycle_hours.sum() * 3600)
T['cost'] = T.cyc * T.tt.map(rate).fillna(0); T['paid_h'] = T.cyc * T.tt.map(hrs_rate).fillna(0)
T['month'] = T.event_date.astype(str).str[:7]
pm = T[T.activity != 'Other'].groupby(['activity', 'person_id', 'month']).agg(cyc=('cyc', 'sum'), ops=('is_op', 'sum'))
pm = pm[pm.ops >= 30]; pm['s_per_op'] = pm.cyc / pm.ops
rows = []
for a in ['Inbound receipt', 'Repalletisation', 'Pallet movement', 'Case picking', 'Piece picking', 'Truck loading', 'Pallet picking']:
    x = T[T.activity == a]; ops = int(x.is_op.sum()); cyc = x.cyc.sum(); cost = x.cost.sum()
    b = pm.loc[a] if a in pm.index.get_level_values(0) else pd.DataFrame()
    actual = cyc / max(ops, 1)
    p25 = float(b.s_per_op.quantile(0.25)) if len(b) >= 5 else np.nan; p50 = float(b.s_per_op.quantile(0.5)) if len(b) >= 5 else np.nan
    rows.append({'activity': a, 'operations': ops, 'scan_work_hours': round(cyc / 3600, 1), 'paid_hours_posted': round(x.paid_h.sum(), 1),
                 'labour_cost_sar_8m': round(cost, 2), 'sar_per_100_ops': round(100 * cost / max(ops, 1), 2), 'seconds_per_op_actual': round(actual, 1),
                 'seconds_per_op_median_person_month': round(p50, 1), 'seconds_per_op_best_quartile': round(p25, 1),
                 'rate_vs_best_quartile_pct': round(100 * (actual / p25 - 1), 1) if p25 == p25 else np.nan,
                 'saving_at_median_rate_sar_yr': round(max(0, 1 - p50 / actual) * cost * 1.5, 0) if p50 == p50 else np.nan,
                 'saving_at_best_quartile_rate_sar_yr': round(max(0, 1 - p25 / actual) * cost * 1.5, 0) if p25 == p25 else np.nan,
                 'person_months_benchmarked': int(len(b)), 'in_deck_slide_6': a != 'Pallet picking'})
cost_act = pd.DataFrame(rows)
print(cost_act[['activity', 'operations', 'sar_per_100_ops', 'rate_vs_best_quartile_pct']].to_string())

# ------------------------------------------------------------------ 2. heads needed vs on shift, by site and hour
sess = pd.read_parquet(f'{SIM}/sim_session_attendance.parquet')
e = sess[sess.eligible_for_benchmark]
tgt = e.groupby('role').span_productivity.quantile([0.5, 0.75]).unstack()
tgt_all = e.span_productivity.quantile([0.5, 0.75])
T['hour'] = T.start_ts.dt.floor('h')
earned = T.groupby(['wh_id', 'hour', 'tran_category']).std_s.sum().div(3600).rename('earned_h').reset_index()
earned['t50'] = earned.tran_category.astype(str).map(tgt[0.5]).fillna(tgt_all[0.5]); earned['t75'] = earned.tran_category.astype(str).map(tgt[0.75]).fillna(tgt_all[0.75])
earned['need50'] = earned.earned_h / earned.t50; earned['need75'] = earned.earned_h / earned.t75
need = earned.groupby(['wh_id', 'hour'])[['earned_h', 'need50', 'need75']].sum()
# on shift: each person session counts for the share of every clock hour its scanned span covers, at its main site
ss = sess[['session_id', 'main_wh', 'session_start', 'session_end']].dropna()
rows = []
for r in ss.itertuples(index=False):
    h0 = r.session_start.floor('h')
    for h in pd.date_range(h0, r.session_end.floor('h'), freq='h'):
        ov = (min(r.session_end, h + pd.Timedelta(hours=1)) - max(r.session_start, h)).total_seconds() / 3600
        if ov > 0: rows.append((r.main_wh, h, ov))
on = pd.DataFrame(rows, columns=['wh_id', 'hour', 'on_shift']).groupby(['wh_id', 'hour']).on_shift.sum()
HH = need.join(on, how='outer').fillna(0).reset_index()
HH['excess50'] = (HH.on_shift - HH.need50).clip(lower=0); HH['excess75'] = (HH.on_shift - HH.need75).clip(lower=0)
HH['date'] = HH.hour.dt.normalize(); HH['weekday'] = HH.hour.dt.day_name(); HH['hod'] = HH.hour.dt.hour
last8 = HH[HH.date > pd.Timestamp('2026-09-30') - pd.Timedelta(weeks=8)]
ndays = last8.groupby(['wh_id', 'weekday']).date.nunique()
prof = last8.groupby(['wh_id', 'weekday', 'hod'])[['on_shift', 'need50', 'need75', 'earned_h']].sum().reset_index()
prof['n_days'] = [ndays.get((w, d), 1) for w, d in zip(prof.wh_id, prof.weekday)]
for c in ('on_shift', 'need50', 'need75', 'earned_h'):
    prof[c] = prof[c] / prof.n_days
hh_kpi = {'excess_head_hours_vs_median_rate': float(HH.excess50.sum() / HH.on_shift.sum()), 'excess_head_hours_vs_top_quartile_rate': float(HH.excess75.sum() / HH.on_shift.sum()),
          'excess_head_hours_vs_median_rate_last8w': float(last8.excess50.sum() / last8.on_shift.sum()), 'excess_head_hours_vs_top_quartile_rate_last8w': float(last8.excess75.sum() / last8.on_shift.sum())}
print({k: round(v, 3) for k, v in hh_kpi.items()})

# ------------------------------------------------------------------ 3. daily workload forecast (earned standard hours per site)
D = T.groupby(['wh_id', T.start_ts.dt.normalize().rename('date')]).std_s.sum().div(3600).rename('earned_h').reset_index()
sites = D.groupby('wh_id').earned_h.sum().sort_values(ascending=False)
sites = sites[sites > 50].index
days = pd.date_range('2026-02-01', '2026-09-30')
D = pd.MultiIndex.from_product([sites, days], names=['wh_id', 'date']).to_frame(index=False).merge(D, how='left').fillna({'earned_h': 0})
W = D.pivot(index='date', columns='wh_id', values='earned_h')


def fc(method, origin, horizon_days):
    hist = W[W.index <= origin]; out = {}
    for d in pd.date_range(origin + pd.Timedelta(days=1), periods=horizon_days):
        same = hist[hist.index.weekday == d.weekday()].tail(4)
        if method == 'last_week': out[d] = same.iloc[-1]
        elif method == 'mean_last4_same_weekday': out[d] = same.mean()
        elif method == 'median_last4_same_weekday': out[d] = same.median()
    return pd.DataFrame(out).T


def score(method, origin, end):
    f = fc(method, origin, (end - origin).days); a = W.loc[f.index]
    return wape(a.to_numpy().ravel(), f.to_numpy().ravel()), f
methods = ['last_week', 'mean_last4_same_weekday', 'median_last4_same_weekday']
VAL_MONTHS = (('2026-05-31', '2026-06-30'), ('2026-06-30', '2026-07-31'), ('2026-07-31', '2026-08-31'))
val = {m: float(np.mean([score(m, pd.Timestamp(o), pd.Timestamp(e))[0] for o, e in VAL_MONTHS])) for m in methods}
best = min(val, key=val.get)
test = {m: score(m, pd.Timestamp('2026-08-31'), pd.Timestamp('2026-09-30'))[0] for m in methods}
_, ftest = score(best, pd.Timestamp('2026-08-31'), pd.Timestamp('2026-09-30'))
foct = fc(best, pd.Timestamp('2026-09-30'), 31)
wf_metrics = {'method_chosen_on_jun_aug': best, **{f'validation_wape_jun_aug::{m}': v for m, v in val.items()}, **{f'test_wape_sep::{m}': v for m, v in test.items()}}
print(wf_metrics)

# ------------------------------------------------------------------ 4. role x activity hours (scan work time)
ra = T.groupby(['wms_role', 'tran_category']).cyc.sum().div(3600).rename('scan_work_hours').reset_index()
ra['share_of_role_hours'] = ra.scan_work_hours / ra.groupby('wms_role').scan_work_hours.transform('sum')
LOW = {'Loading / shipping', 'Picking', 'Internal movement', 'Consolidation', 'Replenishment'}
ra['flag_senior_role_on_floor_work'] = ra.wms_role.isin(['Storekeeper', 'Warehouse manager']) & ra.tran_category.astype(str).isin(LOW)

# ------------------------------------------------------------------ 5. gap to median / top quartile (per full shift, 100% of the gap; annualised)
s2 = sess.copy()
s2['t50'] = s2.role.map(tgt[0.5]); s2['t75'] = s2.role.map(tgt[0.75])
s2['cph'] = s2.labour_cost_base_sar / s2.paid_hours.clip(lower=1e-9)
el = s2[s2.eligible_for_benchmark & s2.t50.notna()]
gap50 = ((el.observed_h - el.earned_h_median_p50 / el.t50).clip(lower=0) * el.cph).sum() * 1.5
gap75 = ((el.observed_h - el.earned_h_median_p50 / el.t75).clip(lower=0) * el.cph).sum() * 1.5
gap_kpi = {'gap_to_median_sar_yr': float(gap50), 'gap_to_top_quartile_sar_yr': float(gap75), 'recoverable_base_sar_yr': float(sess.recoverable_cost_base_sar.sum() * 1.5),
           'gap_to_median_share_of_labour': float(gap50 / (sess.labour_cost_base_sar.sum() * 1.5)), 'gap_to_top_quartile_share_of_labour': float(gap75 / (sess.labour_cost_base_sar.sum() * 1.5))}
print({k: round(v, 3) for k, v in gap_kpi.items()})
pickle.dump({'cost_act': cost_act, 'HH': HH, 'prof': prof, 'hh_kpi': hh_kpi, 'W': W, 'ftest': ftest, 'foct': foct, 'wf_metrics': wf_metrics,
             'ra': ra, 'gap_kpi': gap_kpi, 'targets': tgt, 'targets_all': tgt_all}, open(f'{WORK}/warehouse.pkl', 'wb'))