"""Score combination rules on the forecast race (nested: members picked at 28 Jul, rule scored at 25 Aug;
hold-out = members picked on 28 Jul + 25 Aug, applied once at 22 Sep)."""
import pickle, numpy as np, pandas as pd
from scipy.optimize import linprog

D = pickle.load(open('/home/claude/binzagr/boost/work/fc_race.pkl', 'rb')); FC, ACT = D['FC'], D['ACT']
CUTS = sorted(next(iter(ACT.values())).keys())
C1, C2, C3 = CUTS
SITE = {  # site, raw demand (comparable to this replication): ML weekly WAPE by run, and its two baselines
    'total': dict(daily=0.1627, weekly=0.1618, naive7=0.1618, wd4=0.1618),
    'cluster': dict(daily=0.2133, weekly=0.2327, naive7=0.2532, wd4=0.2291),
    'principal': dict(daily=0.2831, weekly=0.3091, naive7=0.3576, wd4=0.3040),
    'van': dict(daily=0.4043, weekly=0.4110, naive7=0.5330, wd4=0.4390),
    'sku': dict(daily=0.6747, weekly=0.6686, naive7=0.8751, wd4=0.7599)}


def wape(F, A): return np.abs(F - A).sum() / A.sum()


def per_series_abs(lv, m, cuts): return sum(np.abs(FC[lv][m][c] - ACT[lv][c]).sum(1) for c in cuts)


def pooled(lv, m, cuts): return sum(np.abs(FC[lv][m][c] - ACT[lv][c]).sum() for c in cuts) / sum(ACT[lv][c].sum() for c in cuts)


def race(lv, pool, sel, apply, k=3, weights='equal'):
    E = np.stack([per_series_abs(lv, m, sel) for m in pool], 1)            # n x M
    order = np.argsort(E, 1)[:, :k]
    F = np.stack([FC[lv][m][apply] for m in pool], 1)                      # n x M x 4
    out = np.zeros_like(F[:, 0])
    for i in range(F.shape[0]):
        idx = order[i]
        if weights == 'equal': w = np.ones(len(idx))
        else: w = 1 / np.maximum(E[i, idx], 1e-9)
        out[i] = (F[i, idx] * (w / w.sum())[:, None]).sum(0)
    return out


def stack_lp(lv, pool, sel, apply):
    Fm = np.concatenate([np.stack([FC[lv][m][c].ravel() for m in pool], 1) for c in sel])     # rows x M
    a = np.concatenate([ACT[lv][c].ravel() for c in sel]); R, M = Fm.shape
    # minimise sum e  s.t.  e >= Fm w - a, e >= a - Fm w, sum w = 1, w >= 0
    import scipy.sparse as sp
    c = np.r_[np.zeros(M), np.ones(R)]
    I = sp.identity(R, format='csr'); Fs = sp.csr_matrix(Fm)
    A_ub = sp.vstack([sp.hstack([Fs, -I]), sp.hstack([-Fs, -I])]).tocsr(); b_ub = np.r_[a, -a]
    A_eq = sp.csr_matrix(np.r_[np.ones(M), np.zeros(R)][None, :])
    r = linprog(c, A_ub=A_ub, b_ub=b_ub, A_eq=A_eq, b_eq=[1], bounds=[(0, None)] * (M + R), method='highs')
    w = r.x[:M]
    return sum(w[j] * FC[lv][m][apply] for j, m in enumerate(pool)), dict(zip(pool, np.round(w, 3)))


def strategies(lv):
    models = list(FC[lv].keys()); site_pool = [m for m in models if m != 'wd4']
    S = {}
    for m in models: S[f'single:{m}'] = lambda sel, ap, m=m: FC[lv][m][ap]
    S['best single model (picked on validation)'] = lambda sel, ap: FC[lv][min(models, key=lambda m: pooled(lv, m, sel))][ap]
    S['per-series winner'] = lambda sel, ap: race(lv, models, sel, ap, k=1)
    S['site style: per-series top 3, equal'] = lambda sel, ap: race(lv, site_pool, sel, ap, k=3)
    S['per-series top 3, equal, + 4-week avg'] = lambda sel, ap: race(lv, models, sel, ap, k=3)
    S['per-series top 3, inverse-error'] = lambda sel, ap: race(lv, models, sel, ap, k=3, weights='inv')
    S['level top 3 models, equal'] = lambda sel, ap: np.mean([FC[lv][m][ap] for m in sorted(models, key=lambda m: pooled(lv, m, sel))[:3]], 0)
    S['level top 5 models, equal'] = lambda sel, ap: np.mean([FC[lv][m][ap] for m in sorted(models, key=lambda m: pooled(lv, m, sel))[:5]], 0)
    S['mean of all models'] = lambda sel, ap: np.mean([FC[lv][m][ap] for m in models], 0)
    S['median of all models'] = lambda sel, ap: np.median([FC[lv][m][ap] for m in models], 0)
    if lv != 'sku' or True:
        S['stacked weights per level (fit on validation)'] = lambda sel, ap: stack_lp(lv, models, sel, ap)[0]
    return S


if __name__ == '__main__':
    out, picks, weights = [], {}, {}
    for lv in FC:
        S = strategies(lv); rows = []
        for name, fn in S.items():
            v = wape(fn([C1], C2), ACT[lv][C2])                 # nested validation score
            t = wape(fn([C1, C2], C3), ACT[lv][C3])             # hold-out
            rows.append({'level': lv, 'strategy': name, 'val': v, 'holdout': t})
        R = pd.DataFrame(rows)
        combos = R[~R.strategy.str.startswith('single:')]
        pick = combos.loc[combos.val.idxmin()]
        picks[lv] = pick.strategy
        weights[lv] = stack_lp(lv, list(FC[lv].keys()), [C1, C2], C3)[1]
        R['picked'] = R.strategy == pick.strategy
        out.append(R)
        nv, w4 = R.set_index('strategy').loc['single:naive7', 'holdout'], R.set_index('strategy').loc['single:wd4', 'holdout']
        site = R.set_index('strategy').loc['site style: per-series top 3, equal', 'holdout']
        print(f"\n== {lv}: hold-out naive7 {nv:.4f} wd4 {w4:.4f} | site-style replica {site:.4f} | picked on validation: {pick.strategy} -> {pick.holdout:.4f} | site ML daily {SITE[lv]['daily']} weekly {SITE[lv]['weekly']}")
        print(R.sort_values('val').head(10).round(4).to_string(index=False))
    T = pd.concat(out); T.to_csv('/home/claude/binzagr/boost/work/fc_eval.csv', index=False)
    pickle.dump({'table': T, 'picks': picks, 'weights': weights}, open('/home/claude/binzagr/boost/work/fc_eval.pkl', 'wb'))
    print('\nstacked weights fit on 28 Jul + 25 Aug:'); [print(lv, {k: v for k, v in w.items() if v > 0}) for lv, w in weights.items()]