"""Turn the forecast race into the 03_forecast CSVs (backtest chart, metrics, model race, next 12 weeks).

Strategy per level is chosen on the validation weeks only (1-22 Sep), never on the hold-out:
each single model, and "best model per series" (picked on the calibration weeks), are scored on validation;
the lowest weekly error wins. The naive "repeat last week" is the benchmark, not a candidate.
"""
import pickle, numpy as np, pandas as pd
from pcommon import *
import fc_engine as E

Z95 = 1.645
MODELS = ['MA4', 'SES', 'AutoETS', 'AutoTheta', 'Croston', 'CrostonSBA', 'TSB', 'ADIDA', 'IMAPA', 'LightGBM']
UNIT = {'sku': 'cases', 'van': 'SAR', 'principal': 'SAR'}
P = pd.read_parquet(f'{WORK}/products.parquet').set_index('sku')
S = pd.read_parquet(f'{WORK}/sales.parquet')
ORDER_RULE = {'active': 'forecast p50 + safety stock (95% service)',
              'slowing': 'forecast p50 only, no safety stock: no sale for longer than usual',
              'dormant': 'no order: no sale for 8+ weeks, review (delist or promotion?)'}


def selling_status(df):
    """Per series: last week with a sale, weeks since then, usual gap between selling weeks, and a status.
    Same idea as the outlet rule, scaled to each series' own rhythm, so an item that sells every 6 weeks is not dormant after 5:
    dormant = no sale for max(8, 3 x usual gap) weeks; slowing = no sale for max(4, 2 x usual gap) weeks; otherwise active."""
    s = df[df.y > 0].sort_values(['unique_id', 'ds'])
    ids = pd.Index(df.unique_id.unique())
    last = s.groupby('unique_id').ds.max().reindex(ids)
    gap = (s.groupby('unique_id').ds.diff().dt.days / 7).groupby(s.unique_id).median().reindex(ids)
    wks = (DATA_END - last).dt.days // 7
    g = gap.fillna(1.0)
    st = np.where(wks.isna() | (wks >= np.maximum(8, 3 * g)), 'dormant', np.where(wks >= np.maximum(4, 2 * g), 'slowing', 'active'))
    return pd.DataFrame({'last_sale_week': last, 'weeks_since_last_sale': wks, 'usual_gap_weeks': gap.round(1), 'selling_status': st}, index=ids)


def vol_class(tot):
    s = tot.sort_values(ascending=False); c = s.cumsum() / s.sum()
    cls = pd.Series('C', index=s.index); cls[c <= 0.80] = 'A'; cls[(c > 0.80) & (c <= 0.95)] = 'B'
    cls[c[c > 0.80].index[:1]] = 'A'
    return cls


def mae_by_series(d, models):
    return d.groupby('unique_id')[models].apply(lambda g: g.sub(d.loc[g.index, 'y'], axis=0).abs().mean())


def pick(d, winners):
    return d.apply(lambda r: r[winners[r.unique_id]], axis=1)


def run(level):
    R = pickle.load(open(f'{WORK}/race_{level}.pkl', 'rb'))
    df = R['df']
    cv1 = R['cv1'].merge(R['l1'][['unique_id', 'ds', 'LightGBM']], on=['unique_id', 'ds'], how='left')
    cv1 = cv1.merge(E.scale_at(df, sorted(cv1.cutoff.unique())), on=['unique_id', 'cutoff'], how='left')
    tot = df[df.ds <= TRAIN_END].groupby('unique_id').y.sum(); cls = vol_class(tot)
    cv1['cls'] = cv1.unique_id.map(cls)
    cal, val, test = (cv1[cv1.ds.isin(T)].copy() for T in (E.CAL_T, E.VAL_T, E.TEST_T))
    # ---- strategy choice on validation
    cand = {mm: wape(val.y, val[mm]) for mm in MODELS}
    win_cal = mae_by_series(cal, MODELS).idxmin(axis=1)
    cand['Best model per series'] = wape(val.y, pick(val, win_cal))
    strategy = min(cand, key=cand.get)
    win_val = mae_by_series(val, MODELS).idxmin(axis=1)                      # per-series winners (race table + per-series strategy)
    if strategy == 'Best model per series':
        sel_test = win_val; sel_cal = win_cal
    else:
        sel_test = pd.Series(strategy, index=win_val.index); sel_cal = sel_test
    for d, w in ((test, sel_test), (cal, sel_cal), (val, sel_test)):
        d['model'] = d.unique_id.map(w); d['fc'] = pick(d, w)
    # ---- 80% bands from calibration residuals (scaled by the 13-week mean), pooled by volume class
    cal['e'] = (cal.y - cal.fc) / cal.scale
    q = cal.groupby('cls').e.quantile([0.1, 0.9]).unstack()
    test['p10'] = np.minimum((test.fc + test.cls.map(q[0.1]) * test.scale).clip(lower=0), test.fc)
    test['p90'] = np.maximum(test.fc + test.cls.map(q[0.9]) * test.scale, test.fc)
    # ---- promotion-like spikes in the hold-out (week > 5x the series' own 13-week max and > 1,000 units)
    wide = df.pivot(index='unique_id', columns='ds', values='y')
    prior_max = {tw: wide.loc[:, (wide.columns < tw) & (wide.columns >= tw - pd.Timedelta(weeks=13))].max(axis=1) for tw in E.TEST_T}
    test['spike'] = [bool(y > 5 * prior_max[d][u] and y > (1000 if level == 'sku' else 1e12)) for u, d, y in zip(test.unique_id, test.ds, test.y)]
    spike_ids = set(test.loc[test.spike, 'unique_id'])
    # ---- metrics
    m = {'n_series': df.unique_id.nunique(), 'strategy': strategy, 'holdout_wape_model': wape(test.y, test.fc), 'holdout_wape_naive': wape(test.y, test.Naive)}
    m['forecast_value_added'] = 1 - m['holdout_wape_model'] / m['holdout_wape_naive']
    m['band_80_coverage_holdout'] = float(((test.y >= test.p10) & (test.y <= test.p90)).mean())
    m['holdout_bias_p50'] = float(test.fc.sum() / test.y.sum() - 1)
    if spike_ids:
        tx = test[~test.unique_id.isin(spike_ids)]
        m['spike_series_in_holdout'] = len(spike_ids)
        m['holdout_wape_model_excl_spike_series'] = wape(tx.y, tx.fc); m['holdout_wape_naive_excl_spike_series'] = wape(tx.y, tx.Naive)
    m['validation_wape_model'] = wape(val.y, val.fc); m['validation_wape_naive'] = wape(val.y, val.Naive)
    for k, v in cand.items():
        m[f'validation_wape_candidate::{k}'] = v
    for c in ('A', 'B', 'C'):
        t = test[test.cls == c]
        if len(t):
            m[f'holdout_wape_model_class_{c}'] = wape(t.y, t.fc); m[f'holdout_wape_naive_class_{c}'] = wape(t.y, t.Naive)
    # safety stock at the same service level (z = 1.645, one-week cover): sigma of one-week-ahead hold-out errors
    # sigma estimated robustly as 1.25 x mean absolute error (equals the standard deviation for normal errors, less driven by one promotion spike)
    rm = test.groupby('unique_id').apply(lambda g: pd.Series({'rmse_model': 1.25 * (g.y - g.fc).abs().mean(), 'rmse_naive': 1.25 * (g.y - g.Naive).abs().mean()}))
    allw = cv1.copy(); allw['fc'] = pick(allw, sel_test)
    rm['sigma_16w'] = (1.25 * (allw.y - allw.fc).abs().groupby(allw.unique_id).mean()).reindex(rm.index)
    if level == 'sku':
        recent = S[(S.week_end > DATA_END - pd.Timedelta(weeks=13)) & S.sellable]
        ppc = (recent.groupby('sku').net_amt.sum() / recent.groupby('sku').qty_cases.sum()).replace([np.inf], np.nan)
        rm['price'] = rm.index.map(ppc).fillna(0)
    else:
        rm['price'] = 1.0
    m['safety_stock_value_naive_sar'] = Z95 * (rm.rmse_naive * rm.price).sum(); m['safety_stock_value_model_sar'] = Z95 * (rm.rmse_model * rm.price).sum()
    m['safety_stock_change'] = m['safety_stock_value_model_sar'] / m['safety_stock_value_naive_sar'] - 1
    if spike_ids:
        r2 = rm[~rm.index.isin(spike_ids)]
        m['safety_stock_change_excl_spike_series'] = (r2.rmse_model * r2.price).sum() / (r2.rmse_naive * r2.price).sum() - 1
    # ---- forward 12 weeks from 20 Oct with the same strategy
    fw = R['fwd'].merge(R['lf'][['unique_id', 'ds', 'h', 'LightGBM']], on=['unique_id', 'ds'], how='left')
    if strategy == 'Best model per series':
        last8 = cv1[cv1.ds.isin(E.VAL_T.append(E.TEST_T))]
        sel_f = mae_by_series(last8, MODELS).idxmin(axis=1)
    else:
        sel_f = pd.Series(strategy, index=win_val.index)
    fw['model'] = fw.unique_id.map(sel_f); fw['fc'] = pick(fw, sel_f).clip(lower=0)
    fw['scale'] = fw.unique_id.map(E.scale_at(df, [DATA_END]).set_index('unique_id').scale); fw['cls'] = fw.unique_id.map(cls)
    c12 = R['cv12'].merge(R['l12'][['unique_id', 'ds', 'cutoff', 'LightGBM']], on=['unique_id', 'ds', 'cutoff'], how='left')
    c12['h'] = ((c12.ds - c12.cutoff).dt.days // 7).astype(int)
    c12 = c12.merge(E.scale_at(df, sorted(c12.cutoff.unique())), on=['unique_id', 'cutoff'], how='left')
    c12['fc'] = pick(c12, sel_f); c12['cls'] = c12.unique_id.map(cls)
    bins, labs = [0, 1, 4, 8, 12], ['1', '2-4', '5-8', '9-12']
    c12['hb'] = pd.cut(c12.h, bins, labels=labs).astype(str); fw['hb'] = pd.cut(fw.h, bins, labels=labs).astype(str)
    c12['e'] = (c12.y - c12.fc) / c12.scale
    qq = c12.groupby(['cls', 'hb']).e.quantile([0.1, 0.9]).unstack()
    k = list(zip(fw.cls, fw.hb))
    q10 = np.array([qq[0.1].get(x, 0.0) for x in k]); q90 = np.array([qq[0.9].get(x, 0.0) for x in k])
    fw['p10'] = np.minimum((fw.fc + q10 * fw.scale).clip(lower=0), fw.fc); fw['p90'] = np.maximum(fw.fc + q90 * fw.scale, fw.fc)
    return dict(df=df, cv1=cv1, test=test, win_val=win_val, sel_test=sel_test, sel_f=sel_f, cls=cls, tot=tot, m=m, rm=rm, fw=fw, imp=R['imp'],
                strategy=strategy, cand=cand, spike_ids=spike_ids)


def names(level, ids):
    ids = pd.Series(ids, index=ids)
    if level == 'sku':
        return ids.map(P.prd_desc_en).fillna(ids.map(S.drop_duplicates('sku').set_index('sku').prd_desc))
    return ids


if __name__ == '__main__':
    out = {lv: run(lv) for lv in ('sku', 'van', 'principal')}
    pickle.dump(out, open(f'{WORK}/forecast_results.pkl', 'wb'))
    series_rows, bt_rows, fw_rows, met_rows, race_rows = [], [], [], [], []
    for lv, o in out.items():
        ids = o['df'].unique_id.unique()
        tst = o['test'].groupby('unique_id').apply(lambda g: pd.Series({'holdout_actual': round(g.y.sum(), 2), 'holdout_wape_model': round(wape(g.y, g.fc), 4),
                                                                        'holdout_wape_naive': round(wape(g.y, g.Naive), 4), 'holdout_spike': bool(g.spike.any())}))
        sr = pd.DataFrame({'level': lv, 'series_id': ids, 'series_name': names(lv, ids).reindex(ids).values, 'unit': UNIT[lv],
                           'volume_class': pd.Series(ids).map(o['cls']).values, 'training_total': pd.Series(ids).map(o['tot']).round(2).values,
                           'model_used': pd.Series(ids).map(o['sel_test']).values, 'race_winner_validation': pd.Series(ids).map(o['win_val']).values})
        if lv == 'sku':
            sr['principal'] = pd.Series(ids).map(P.l1_desc).values; sr['units_per_case'] = pd.Series(ids).map(P.case_factor).values
        sst = selling_status(o['df'])
        sr = sr.merge(sst, left_on='series_id', right_index=True, how='left')
        series_rows.append(sr.merge(tst, left_on='series_id', right_index=True, how='left'))
        h = o['df'].copy(); h['split'] = np.where(h.ds.isin(E.TEST_T), 'holdout', 'history')
        h = h.merge(o['test'][['unique_id', 'ds', 'fc', 'p10', 'p90', 'Naive', 'model']], on=['unique_id', 'ds'], how='left')
        bt_rows.append(pd.DataFrame({'level': lv, 'series_id': h.unique_id, 'week_start': h.ds - pd.Timedelta(days=6), 'week_end': h.ds, 'split': h.split,
                                     'actual': h.y.round(3), 'forecast_p10': h.p10.round(3), 'forecast_p50': h.fc.round(3), 'forecast_p90': h.p90.round(3),
                                     'naive_repeat_last_week': h.Naive.round(3), 'model': h.model, 'unit': UNIT[lv]}))
        f = o['fw']
        fr = pd.DataFrame({'level': lv, 'series_id': f.unique_id, 'week_start': f.ds - pd.Timedelta(days=6), 'week_end': f.ds, 'horizon_weeks': f.h,
                           'forecast_p10': f.p10.round(3), 'forecast_p50': f.fc.round(3), 'forecast_p90': f.p90.round(3), 'model': f.model, 'unit': UNIT[lv]})
        fr['selling_status'] = fr.series_id.map(sst.selling_status)
        if lv == 'sku':
            # no safety stock for an item that has stopped selling, and no order for a dormant one (it goes to review instead)
            ss = fr.series_id.map(Z95 * o['rm'].sigma_16w).fillna(0)
            fr['safety_stock_cases'] = np.where(fr.selling_status == 'active', ss, 0.0).round(2)
            fr['suggested_order_qty_cases'] = np.where(fr.selling_status == 'dormant', 0.0, np.ceil(fr.forecast_p50 + fr.safety_stock_cases))
            fr['order_rule'] = fr.selling_status.map(ORDER_RULE)
        fw_rows.append(fr)
        for k, v in o['m'].items():
            met_rows.append({'level': lv, 'metric': k, 'value': v if isinstance(v, str) else round(float(v), 4)})
        holdout_best = mae_by_series(o['test'], MODELS).idxmin(axis=1)
        for mm in MODELS:
            race_rows.append({'level': lv, 'model': mm, 'family': E.FAMILY[mm], 'status': 'used for every series' if o['strategy'] == mm else 'run',
                              'series_won_validation': int((o['win_val'] == mm).sum()), 'series_won_holdout': int((holdout_best == mm).sum()),
                              'share_of_volume_won_holdout': round(o['tot'].reindex(holdout_best[holdout_best == mm].index).sum() / o['tot'].sum(), 4),
                              'validation_wape_if_used_for_all': round(o['cand'][mm], 4), 'holdout_wape_if_used_for_all': round(wape(o['test'].y, o['test'][mm]), 4)})
        race_rows.append({'level': lv, 'model': 'Best model per series', 'family': 'Selection', 'status': 'used' if o['strategy'] == 'Best model per series' else 'tested, not better on validation',
                          'validation_wape_if_used_for_all': round(o['cand']['Best model per series'], 4)})
        race_rows.append({'level': lv, 'model': 'Naive (repeat last week)', 'family': 'Baseline', 'status': 'benchmark',
                          'validation_wape_if_used_for_all': round(o['m']['validation_wape_naive'], 4), 'holdout_wape_if_used_for_all': round(o['m']['holdout_wape_naive'], 4)})
        for fm in ('Chronos-2', 'TimesFM 2.5', 'TiRex'):
            race_rows.append({'level': lv, 'model': fm, 'family': 'Foundation', 'status': 'not run: model weights could not be downloaded in the build environment'})
    write_csv(pd.concat(series_rows), '03_forecast', 'forecast_series.csv')
    write_csv(pd.concat(bt_rows), '03_forecast', 'forecast_backtest_weekly.csv')
    write_csv(pd.concat(fw_rows), '03_forecast', 'forecast_next_12_weeks.csv')
    write_csv(pd.DataFrame(met_rows), '03_forecast', 'forecast_metrics.csv')
    write_csv(pd.DataFrame(race_rows), '03_forecast', 'forecast_model_race.csv')
    for lv, o in out.items():
        print(lv, {k: (round(v, 3) if isinstance(v, float) else v) for k, v in o['m'].items() if 'candidate' not in k}, 'spikes:', sorted(o['spike_ids'])[:6])