"""Stage 2: promotions (F08), sell-in for all channels (F02), targets (F01)."""
import numpy as np, pandas as pd
from common import *
import warnings; warnings.filterwarnings('ignore')

r = rng(2)
S = real_sales(); prod = load('dim_product_ext'); outl = load('dim_outlet'); acc = load('dim_account')
cal = calendar(); days = cal.date
twin = outl[outl.channel_id == 'GR'][['outlet_id', 'source_codes']].assign(code=lambda d: d.source_codes.str.split('|')).explode('code')
code2main = dict(zip(twin.code, twin.outlet_id))
base = prod[prod.channel_variant_of.isna()].set_index('sku')

# ---------------- F08 promotions (MT, DS, EC accounts)
pacc = acc[acc.channel_id.isin(['MT', 'DS', 'EC'])]
nprom = {'MT': 75, 'DS': 40, 'EC': 30}
w_sku = base.real_gr_value_sar ** 0.7; w_sku = w_sku / w_sku.sum()
rows, pid = [], 0
ramadan_end = pd.Timestamp('2026-03-19'); bts = pd.Timestamp('2026-08-20')
for _, a in pacc.iterrows():
    n = int(nprom[a.channel_id] * {'A': 1.2, 'B': 1.0, 'C': 0.7}[a.account_tier])
    for k in range(n):
        pid += 1
        u = r.random()
        start = (START + pd.Timedelta(days=int(r.integers(0, 18))) if u < 0.2 else
                 bts + pd.Timedelta(days=int(r.integers(-10, 25))) if u < 0.4 else
                 START + pd.Timedelta(days=int(r.integers(0, (END - START).days - 10))))
        dur = int(r.choice([7, 10, 14, 21], p=[0.3, 0.2, 0.35, 0.15]))
        mech = r.choice(['price-off', 'multibuy', 'display', 'free goods'], p=[0.5, 0.2, 0.15, 0.15])
        disc = round(float(r.choice([0.10, 0.15, 0.20, 0.25, 0.30], p=[0.2, 0.3, 0.25, 0.15, 0.1])), 2)
        nsku = int(r.choice([1, 2, 3, 4], p=[0.4, 0.3, 0.2, 0.1]))
        skus = r.choice(w_sku.index, nsku, replace=False, p=w_sku.values)
        planned = round(float(r.uniform(0.30, 0.70)), 2)                         # plans are optimistic
        true = float(np.clip(r.lognormal(np.log(0.22), 0.6) * (1 + disc), 0, 1.2))  # realised uplift, hidden truth
        if r.random() < 0.15: true *= 0.15                                        # some promotions barely work
        display = mech == 'display' or r.random() < 0.3
        for s in skus:
            rows.append((f'PR{pid:04d}', a.account_id, a.channel_id, s, start.date(), (start + pd.Timedelta(days=dur - 1)).date(), mech, disc,
                         planned, round(true, 3), display))
PR = pd.DataFrame(rows, columns=['promotion_id', 'account_id', 'channel_id', 'sku', 'start_date', 'end_date', 'mechanic', 'discount_pct',
                                 'planned_uplift_pct', 'sim_true_uplift_pct', 'display_flag'])
PR['end_date'] = PR.end_date.clip(upper=END.date())
print('  promotions', PR.promotion_id.nunique(), 'lines', len(PR))

# ---------------- F02 sell-in
gr = S.copy()
gr['outlet_id'] = gr.outlet_id.map(code2main).fillna(gr.outlet_id)
gr['channel_id'] = 'GR'; gr['account_id'] = 'GR-' + gr.outlet_id
gr['list_price_sar'] = gr.prd_cd.map(base.list_price_sar)
gr['discount_sar'] = np.where(gr.is_foc, 0, (gr.list_price_sar * gr.qty_base_ea - gr.net_amt).clip(lower=0)).round(2)
gr['promotion_id'] = None; gr['po_id'] = None; gr['ship_to_city'] = 'Riyadh'; gr['anomaly_label'] = None; gr['basis'] = 'Actual'
gr['sku'] = gr.prd_cd
grd = gr[~gr.is_foc].groupby('doc_date').net_amt.sum().reindex(days, fill_value=0)
smooth = grd.rolling(7, center=True, min_periods=1).mean()      # real daily seasonality (incl. Ramadan/Eid, weekly rhythm smoothed)

lines_lambda = {'MT': {'XL': 45, 'L': 35, 'M': 25, 'S': 16, 'XS': 12}, 'DS': 18, 'WS': 26, 'EC': 40}
cyc = {'MT': 2, 'DS': 1, 'WS': 2, 'EC': 6}
tilt = {'MT': {'Unilever': 1.3, 'Non-Food': 1.3, 'Chilled': 1.1}, 'DS': {'Food': 1.2, 'Drinks': 1.2}, 'WS': {'Drinks': 1.3, 'Food': 1.2, 'Ice Cream': 0.3, 'Chilled': 0.4},
        'EC': {'Non-Food': 1.6, 'Unilever': 1.4, 'Ice Cream': 0.1, 'Chilled': 0.3}}
wd = ['Saturday', 'Sunday', 'Monday', 'Tuesday', 'Wednesday', 'Thursday']
gr_total = grd.sum(); share = P['channel_share']
parts = []
for chn in ['MT', 'DS', 'WS', 'EC']:
    st = outl[outl.channel_id == chn]
    w = base.real_gr_value_sar ** (0.9 if chn == 'WS' else 0.75) * base.division.map(tilt[chn]).fillna(1.0)
    p = (w / w.sum()).values; skus = w.index.values
    ev = []
    for _, o in st.iterrows():
        od = r.choice(wd, cyc[chn], replace=False) if cyc[chn] < 6 else wd
        dd = days[days.dt.day_name().isin(od)]
        dd = dd[(dd >= pd.Timestamp(o.open_date)) & ((pd.isna(o.close_date)) | (dd <= pd.Timestamp(o.close_date) if pd.notna(o.close_date) else True))]
        for d in dd: ev.append((o.outlet_id, o.account_id, o.city, o.size_band, d))
    E = pd.DataFrame(ev, columns=['outlet_id', 'account_id', 'city', 'size_band', 'doc_date'])
    lam = E.size_band.map(lines_lambda['MT']) if chn == 'MT' else pd.Series(lines_lambda[chn], index=E.index)
    nl = np.maximum(3, r.poisson(lam.values))
    E['order_no'] = [f'{chn}{i+1:06d}' for i in range(len(E))]
    L = E.loc[E.index.repeat(nl)].reset_index(drop=True)
    L['doc_line_seq'] = L.groupby('order_no').cumcount() + 1
    L['sku'] = r.choice(skus, len(L), p=p)
    L = L.drop_duplicates(['order_no', 'sku']).reset_index(drop=True)
    if chn == 'DS':   # discounters take the discounter packs where they exist
        var = prod[prod.channel_variant_of.notna()].set_index('channel_variant_of').sku
        m = L.sku.isin(var.index) & (r.random(len(L)) < 0.7)
        L.loc[m, 'sku'] = L.loc[m, 'sku'].map(var)
    L['raw'] = r.lognormal(0, 0.9, len(L))
    # promotions: uplift volume and discount price when the line falls in a promo window
    pr = PR[PR.account_id.isin(L.account_id.unique())][['promotion_id', 'account_id', 'sku', 'start_date', 'end_date', 'discount_pct', 'sim_true_uplift_pct']]
    L['_key'] = L.account_id + '|' + L.sku.astype(str)
    pr = pr.assign(_key=pr.account_id + '|' + pr.sku.astype(str))
    M = L[['_key', 'doc_date']].reset_index().merge(pr, on='_key')
    M = M[(M.doc_date >= pd.to_datetime(M.start_date)) & (M.doc_date <= pd.to_datetime(M.end_date))].drop_duplicates('index')
    L['promotion_id'] = None; L['promo_disc'] = 0.0; L['uplift'] = 0.0
    L.loc[M['index'], 'promotion_id'] = M.promotion_id.values
    L.loc[M['index'], 'promo_disc'] = M.discount_pct.values
    L.loc[M['index'], 'uplift'] = M.sim_true_uplift_pct.values
    L['raw'] *= (1 + L.uplift)
    # scale raw to the channel's daily value target (real grocery seasonality × channel share)
    tgt = smooth * share[chn] / share['GR']
    day_raw = L.groupby('doc_date').raw.sum()
    L['value'] = L.raw * L.doc_date.map(tgt / day_raw)
    L['value'] *= tgt.sum() / L.value.sum()   # Friday value spread over order days
    lp = L.sku.map(prod.set_index('sku').list_price_sar)
    cs = L.sku.map(prod.set_index('sku').case_size)
    L['list_price_sar'] = lp
    L['unit_net_price_per_ea'] = (lp * (1 - P['channel_discount'][chn]) * r.normal(1, 0.025, len(L)) * (1 - L.promo_disc)).round(4)
    xq = (L.value / (L.unit_net_price_per_ea * cs)).values
    L['qty'] = (np.floor(xq) + (r.random(len(L)) < (xq - np.floor(xq)))).astype(int)   # unbiased rounding keeps the value target
    L = L[L.qty > 0].reset_index(drop=True); cs = L.sku.map(prod.set_index('sku').case_size); lp = L.list_price_sar
    L['doc_line_seq'] = L.groupby('order_no').cumcount() + 1
    L['conv_factor'] = cs; L['qty_base_ea'] = L.qty * cs
    L['net_amt'] = (L.qty_base_ea * L.unit_net_price_per_ea).round(2)
    L['discount_sar'] = ((lp - L.unit_net_price_per_ea) * L.qty_base_ea).round(2)
    L['channel_id'] = chn; L['doc_no'] = L.order_no; L['po_id'] = np.where(chn in ('MT', 'DS', 'EC'), 'PO-' + L.order_no, None)
    L['ship_to_city'] = L.city; L['anomaly_label'] = None
    parts.append(L)
    print(f'  {chn}: orders {len(E):,} lines {len(L):,} value SAR {L.net_amt.sum()/1e6:.2f}M (target {tgt.sum()/1e6:.2f}M)')
X = pd.concat(parts, ignore_index=True)

# ---------------- anomalies for task 9 (and 8): wholesale mostly
ws = X.index[X.channel_id == 'WS']
big = r.choice(ws, 60, replace=False)
X.loc[big, 'qty'] = (X.loc[big, 'qty'] * r.integers(6, 11, 60)).astype(int); X.loc[big, 'anomaly_label'] = 'large_irregular'
par = r.choice(np.setdiff1d(ws, big), 40, replace=False)
X.loc[par, 'qty'] = (X.loc[par, 'qty'] * r.integers(3, 6, 40)).astype(int)
X.loc[par, 'ship_to_city'] = r.choice(['Jazan', 'Abha', 'Tabuk', 'Hail'], 40); X.loc[par, 'anomaly_label'] = 'parallel_trade'
cheap = r.choice(X.index[X.channel_id.isin(['WS', 'DS']) & X.anomaly_label.isna()], 90, replace=False)
X.loc[cheap, 'unit_net_price_per_ea'] = (X.loc[cheap, 'list_price_sar'] * r.uniform(0.62, 0.72, 90)).round(4); X.loc[cheap, 'anomaly_label'] = 'price_below_floor'
X['qty_base_ea'] = X.qty * X.conv_factor; X['net_amt'] = (X.qty_base_ea * X.unit_net_price_per_ea).round(2)
X['discount_sar'] = ((X.list_price_sar - X.unit_net_price_per_ea) * X.qty_base_ea).round(2)

X['dist_cd'] = 'SIM'; X['cust_cd'] = X.outlet_id; X['slsman_cd'] = X.account_id.map(acc.set_index('account_id').kam_employee_id)
X['prd_cd'] = X.sku; X['uom_cd'] = 'CS'; X['is_foc'] = False; X['basis'] = 'Simulated'
cols = ['channel_id', 'account_id', 'outlet_id', 'dist_cd', 'doc_no', 'doc_line_seq', 'doc_date', 'cust_cd', 'slsman_cd', 'sku', 'uom_cd', 'qty',
        'conv_factor', 'qty_base_ea', 'is_foc', 'net_amt', 'unit_net_price_per_ea', 'list_price_sar', 'discount_sar', 'promotion_id', 'po_id',
        'ship_to_city', 'anomaly_label', 'basis']
FI = pd.concat([gr[cols], X[cols]], ignore_index=True)
FI['doc_date'] = FI.doc_date.dt.date
FI['sku'] = FI.sku.astype(str); FI['qty'] = FI.qty.astype(float)
save(FI, 'fact_sell_in')
print('  channel mix (value):', (FI[~FI.is_foc].groupby('channel_id').net_amt.sum() / FI[~FI.is_foc].net_amt.sum()).round(3).to_dict())

# promotions: add budget and planned volume from baseline sell-in
pv = X.dropna(subset=['promotion_id']).groupby('promotion_id').agg(promo_units=('qty_base_ea', 'sum'), promo_value=('net_amt', 'sum'))
PR = PR.merge(pv, left_on='promotion_id', right_index=True, how='left')
nsku = PR.groupby('promotion_id').sku.transform('count')
PR['planned_volume_units'] = (PR.promo_units.fillna(0) / nsku / (1 + PR.sim_true_uplift_pct) * (1 + PR.planned_uplift_pct)).round()
PR['budget_sar'] = (PR.promo_value.fillna(0) / nsku * PR.discount_pct * 1.1 + np.where(PR.display_flag, 2500, 0)).round(-1)
PR = PR.drop(columns=['promo_units', 'promo_value'])
save(PR, 'fact_promotion')

# ---------------- F01 targets: monthly by channel × account/route × brand, and daily by channel × key SKU
fi = FI[~FI.is_foc].copy()
fi['month'] = pd.to_datetime(fi.doc_date).dt.to_period('M').dt.to_timestamp()
fi['brand'] = fi.sku.map(prod.set_index('sku').brand).fillna('Other')
fi['division'] = fi.sku.map(prod.set_index('sku').division)
fi['unit'] = np.where(fi.channel_id == 'GR', 'R-' + fi.groupby('outlet_id').outlet_id.transform('first').map(outl.set_index('outlet_id').main_van).fillna('NA').astype(str), fi.account_id)
fi.loc[fi.channel_id == 'GR', 'unit'] = 'R-' + S.slsman_cd.astype(str).reindex(fi.index[fi.channel_id == 'GR']).values
mo = fi.groupby(['channel_id', 'unit', 'division', 'brand', 'month']).agg(v=('net_amt', 'sum'), u=('qty_base_ea', 'sum')).reset_index()
base_run = mo.groupby(['channel_id', 'unit', 'division', 'brand']).v.transform('mean')   # plan anchored on the period run rate
mi = mo.month.map(mo.groupby('month').v.sum() / mo.groupby('month').v.sum().mean())     # seasonality index from the real months
g = mo.division.map(P['growth_target_by_division'])
mo['target_value_sar'] = (base_run * mi * (1 + g) * r.normal(1.0, 0.08, len(mo))).round(0)
mo['target_units'] = (mo.target_value_sar / (mo.v / mo.u.replace(0, np.nan))).round(0)
mo['planned_margin_sar'] = (mo.target_value_sar * mo.division.map(P['division_margin'])).round(0)
mo['plan_version'] = np.where(mo.month < pd.Timestamp('2026-07-01'), 'Budget 2026', 'Re-forecast Jul 2026')
mo = mo.rename(columns={'unit': 'account_or_route_id', 'month': 'period_month'})
mo['period_month'] = mo.period_month.dt.date
mo.insert(0, 'target_id', [f'TG{i+1:06d}' for i in range(len(mo))])
save(mo.drop(columns=['v', 'u']), 'fact_sales_target')
# daily: channel × key SKU (top 50 by value) phased by selling days, Friday weight 0.3
keys = fi.groupby('sku').net_amt.sum().nlargest(50).index
dm = fi[fi.sku.isin(keys)].groupby(['channel_id', 'sku', 'month']).net_amt.sum().reset_index()
dm['target_month'] = (dm.net_amt * (1 + dm.sku.map(prod.set_index('sku').division).map(P['growth_target_by_division'])) * r.normal(1, 0.06, len(dm)))
cal['month'] = cal.date.dt.to_period('M').dt.to_timestamp(); cal['w'] = np.where(cal.date.dt.day_name() == 'Friday', 0.3, 1.0)
cal['wshare'] = cal.w / cal.groupby('month').w.transform('sum')
dd = dm.merge(cal[['date', 'month', 'wshare']], on='month')
dd['target_value_sar'] = (dd.target_month * dd.wshare).round(2)
dd = dd.rename(columns={'date': 'target_date'})[['target_date', 'channel_id', 'sku', 'target_value_sar']]
dd['target_date'] = dd.target_date.dt.date
save(dd, 'fact_sales_target_daily')