"""Outlet master for the outlet and van screens: status, spend band, churn risk, lapsed value, nearest active van, last SKUs, offers."""
import numpy as np, pandas as pd, pickle
from sklearn.neighbors import BallTree
from pcommon import *

S = pd.read_parquet(f'{WORK}/sales.parquet'); C = pd.read_parquet(f'{CLEAN}/dim_customer.parquet'); P = pd.read_parquet(f'{WORK}/products.parquet').set_index('sku')
C['outlet_id'] = outlet_id(C.dist_cd, C.cust_cd); C = C.set_index('outlet_id')
chmap = pd.read_parquet(f'{SIM}/inf_customer_channel_map.parquet'); chmap = chmap[chmap.level == 'cust_hier2'].set_index('code').inferred_channel
churn = pickle.load(open(f'{WORK}/churn.pkl', 'rb'))['now'].set_index('outlet_id')
reco = pd.read_parquet(f'{WORK}/reco_final.parquet')
B = S[S.sellable & (S.net_amt > 0)]
g = B.groupby('outlet_id')
O = pd.DataFrame({'first_purchase': g.date.min(), 'last_purchase': g.date.max(), 'purchase_days': g.date.nunique(), 'invoices': g.doc_no.nunique(), 'sales_window_sar': g.net_amt.sum()})
gaps = B.groupby('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)
O['usual_gap_days'] = gaps
O['days_since_last'] = (DATA_END - O.last_purchase).dt.days
lim = np.maximum(60, 3 * O.usual_gap_days.fillna(60))
O['status'] = np.where(O.days_since_last > lim, 'lapsed', np.where(O.days_since_last > 30, 'slowing', 'active'))
# annual value: spend in the 26 weeks up to the last purchase (or since the first purchase if shorter), annualised
def annual(o):
    lo = max(o.first_purchase, o.last_purchase - pd.Timedelta(days=182))
    days = max((o.last_purchase - lo).days + (o.usual_gap_days if o.usual_gap_days == o.usual_gap_days else 30), 30)
    return o.sales_26w_before_last * 365 / days
B2 = B.merge(O[['last_purchase']], left_on='outlet_id', right_index=True)
O['sales_26w_before_last'] = B2[B2.date > B2.last_purchase - pd.Timedelta(days=182)].groupby('outlet_id').net_amt.sum()
O['annual_value_sar'] = O.sales_window_sar * 365 / np.maximum((DATA_END - O.first_purchase).dt.days + 1, 56)   # average per year since the first purchase (conservative)
O['annual_run_rate_while_active_sar'] = O.apply(annual, axis=1)
O['spend_band'] = pd.cut(O.annual_value_sar.rank(pct=True), [0, 0.4, 0.7, 0.9, 1.0], labels=['D (bottom 40%)', 'C (next 30%)', 'B (next 20%)', 'A (top 10%)']).astype(str)
r13 = B[B.date > DATA_END - pd.Timedelta(days=91)]
O['main_van'] = r13.groupby('outlet_id').van.agg(lambda x: x.value_counts().index[0])
O['main_van'] = O.main_van.fillna(B.sort_values('date').groupby('outlet_id').van.last())
O['vans_in_window'] = g.van.nunique()
O['last_van'] = B.sort_values('date').groupby('outlet_id').van.last()
O['churn_risk_8w'] = churn.p.reindex(O.index)
O['churn_risk_tier'] = np.select([O.churn_risk_8w >= 0.25, O.churn_risk_8w >= 0.10, O.churn_risk_8w.notna()], ['High', 'Medium', 'Low'], '')
O.loc[O.churn_risk_tier == '', 'churn_risk_tier'] = 'Not scored (no purchase in 13 weeks)'
O.loc[O.status == 'lapsed', 'churn_risk_tier'] = 'Lapsed'
# last SKUs (most recent purchase day, top 3 by value) and offers
lastday = B.merge(O[['last_purchase']], left_on='outlet_id', right_index=True)
lastday = lastday[lastday.date == lastday.last_purchase].groupby(['outlet_id', 'sku']).net_amt.sum().reset_index().sort_values(['outlet_id', 'net_amt'], ascending=[True, False])
lastday['n'] = lastday.groupby('outlet_id').cumcount()
for i in range(3):
    O[f'last_sku_{i + 1}'] = lastday[lastday.n == i].set_index('outlet_id').sku.map(P.prd_desc_en).reindex(O.index)
for i in range(1, 6):
    r = reco[reco['rank'] == i].set_index('outlet_id')
    O[f'offer_{i}'] = r.sku_name.reindex(O.index); O[f'offer_{i}_sku'] = r.sku.reindex(O.index); O[f'offer_{i}_value_sar'] = r.expected_value_8w_sar.reindex(O.index)
# nearest active van for lapsed outlets (nearest outlet that bought in the last 28 days, by straight-line distance)
cc = C.reindex(O.index)
O['lat'] = cc.latitude.values; O['lon'] = cc.longitude.values
act = O[(O.days_since_last <= 28) & O.lat.notna() & O.lon.notna()]
tree = BallTree(np.radians(act[['lat', 'lon']].to_numpy()), metric='haversine')
need = O[O.lat.notna() & O.lon.notna()]
d, ix = tree.query(np.radians(need[['lat', 'lon']].to_numpy()), k=2)
self_hit = act.index.to_numpy()[ix[:, 0]] == need.index.to_numpy()
pick = np.where(self_hit, ix[:, 1], ix[:, 0]); dist = np.where(self_hit, d[:, 1], d[:, 0]) * 6371.0
O.loc[need.index, 'nearest_active_outlet'] = act.index.to_numpy()[pick]
O.loc[need.index, 'nearest_active_van'] = act.main_van.to_numpy()[pick]
O.loc[need.index, 'nearest_active_km'] = np.round(dist, 2)
O['outlet_name'] = cc.cust_name.values; O['outlet_name_ar'] = cc.cust_name2.values
O['customer_type'] = np.where(cc.cust_type.values == 'Q', 'Credit', 'Cash'); O['payment_term'] = cc.invterm_cd.values
O['balance_sar'] = cc.outstanding_bal.values; O['channel_code'] = cc.cust_hier2.values
O['channel_inferred'] = pd.Series(cc.cust_hier2.astype(str).values, index=O.index).map(chmap).fillna('unknown')
lapsed = O[O.status == 'lapsed']
kpi = {'buying_outlets': len(O), 'active': int((O.status == 'active').sum()), 'slowing': int((O.status == 'slowing').sum()), 'lapsed': int(len(lapsed)),
       'lapsed_annual_value_sar': float(lapsed.annual_value_sar.sum()), 'lapsed_run_rate_value_sar': float(lapsed.annual_run_rate_while_active_sar.sum()), 'high_risk_outlets': int((O.churn_risk_tier == 'High').sum()),
       'medium_risk_outlets': int((O.churn_risk_tier == 'Medium').sum()), 'low_risk_outlets': int((O.churn_risk_tier == 'Low').sum()),
       'outlets_2plus_vans': float((O.vans_in_window >= 2).mean()), 'master_outlets_with_gps': int(C.latitude.notna().sum())}
print(kpi)
O.to_parquet(f'{WORK}/outlets.parquet'); pickle.dump({'kpi': kpi}, open(f'{WORK}/outlets_kpi.pkl', 'wb'))
print(O.churn_risk_tier.value_counts().to_dict())