"""Replicate the site's forecast test and compare model-combination rules.

Site protocol (Forecast proof page): daily net van sales (SAR), trained to 22 Sep 2026, judged on 23 Sep - 20 Oct
(28 days, fixed origin), error = weekly WAPE pooled over series and the 4 weeks; baselines = last week repeated
(naive7) and the 4-week weekday average (wd4). Levels: company total, 6 clusters, 16 principals, 92 vans, 548 SKUs.

Here: the same data, series and protocol. Candidate pool = statistical models (daily and weekly) + LightGBM (daily
and weekly direct models) + the two baselines. The foundation models the site also uses (Chronos-2, TimesFM 2.5,
TiRex) cannot be downloaded in this environment, so they are not in the pool.

Origins: 28 Jul and 25 Aug (validation) and 22 Sep (hold-out). Combination rules that pick members use 28 Jul to
pick and are scored on 25 Aug (nested), so every rule is chosen on validation only; the hold-out is used once.
"""
import json, pickle, time, warnings
import numpy as np, pandas as pd
from scipy.optimize import linprog
warnings.filterwarnings('ignore')
from statsforecast import StatsForecast
from statsforecast.models import (AutoETS, AutoTheta, ADIDA, IMAPA, TSB, CrostonSBA,
                                  SimpleExponentialSmoothingOptimized)
import lightgbm as lgb

ROOT = '/home/claude/binzagr'
OUT = f'{ROOT}/boost/work'
import os; os.makedirs(OUT, exist_ok=True)
s = pd.read_parquet(f'{ROOT}/pred/work/sales.parquet')
P = pd.read_parquet(f'{ROOT}/pred/work/products.parquet').set_index('sku')
V = json.load(open(f'{ROOT}/boost/site_vans.json')); VMAP = dict(V['all']); FCV = [v[0] for v in V['fc']]
s['cluster'] = s.van.map(VMAP); s['principal'] = s.sku.map(P.l1_desc)
DATES = pd.date_range('2026-03-01', '2026-10-20')
CUTS = [pd.Timestamp('2026-07-28'), pd.Timestamp('2026-08-25'), pd.Timestamp('2026-09-22')]
H = 28
HOLI = set(pd.date_range('2026-03-19', '2026-03-23')) | set(pd.date_range('2026-05-25', '2026-05-30')) | {pd.Timestamp('2026-09-23')}


def panel(key, ids=None):
    d = s.groupby([key, 'date']).net_amt.sum().unstack(fill_value=0).reindex(columns=DATES, fill_value=0)
    return d if ids is None else d.reindex(ids, fill_value=0)


LEVELS = {'total': s.groupby('date').net_amt.sum().reindex(DATES, fill_value=0).to_frame('Company total').T,
          'cluster': panel('cluster')}
pr = panel('principal'); LEVELS['principal'] = pr.loc[pr.loc[:, :CUTS[-1]].sum(1).sort_values(ascending=False).index[:16]]
LEVELS['van'] = panel('van', FCV)
sk = panel('sku'); rows = s[s.date <= CUTS[-1]].groupby('sku').size()
LEVELS['sku'] = sk.loc[rows.sort_values(ascending=False).index[:548]]
# grouping used as a categorical by the global models
sku_cluster = s.groupby('sku').cluster.agg(lambda x: x.value_counts().index[0])
GROUP = {'total': None, 'cluster': None, 'principal': None,
         'van': pd.Series(FCV).map(VMAP).values, 'sku': LEVELS['sku'].index.map(sku_cluster).values}


def weekly_actual(Y, cut):
    i = DATES.get_loc(cut)
    return Y.values[:, i + 1:i + 1 + H].reshape(len(Y), 4, 7).sum(2)


# ------------------------------------------------------------------ baselines
def baselines(Y, cut):
    i = DATES.get_loc(cut); a = Y.values
    last = a[:, i - 6:i + 1]
    wd = np.stack([a[:, i - 6 - 7 * k:i + 1 - 7 * k] for k in range(4)]).mean(0)
    to_w = lambda d: np.tile(d, 4).reshape(len(Y), 4, 7).sum(2)
    return {'naive7': to_w(last), 'wd4': to_w(wd)}


# ------------------------------------------------------------------ statistical models
def sf_daily(Y, cut):
    df = Y.loc[:, :cut].stack().rename('y').reset_index(); df.columns = ['unique_id', 'ds', 'y']
    models = [AutoETS(season_length=7, alias='ETS_d'), AutoTheta(season_length=7, alias='Theta_d'), ADIDA(alias='ADIDA_d'),
              IMAPA(alias='IMAPA_d'), TSB(alpha_d=0.2, alpha_p=0.2, alias='TSB_d'), CrostonSBA(alias='SBA_d')]
    fc = StatsForecast(models=models, freq='D', n_jobs=2).forecast(df=df, h=H)
    if 'unique_id' not in fc.columns: fc = fc.reset_index()
    out = {}
    for m in [x.alias for x in models]:
        M = fc.pivot(index='unique_id', columns='ds', values=m).reindex(Y.index).values.clip(min=0)
        out[m] = M.reshape(len(Y), 4, 7).sum(2)
    return out


def to_weekly(Y):
    """Wed-Tue weeks ending on Tuesdays (the cut-offs are Tuesdays); drops the partial first week."""
    i0 = (DATES.get_loc(pd.Timestamp('2026-03-10')) - 6)          # first full week: 4-10 Mar
    a = Y.values[:, i0:]; n = a.shape[1] // 7
    ends = DATES[i0 + 6::7][:n]
    return pd.DataFrame(a[:, :n * 7].reshape(len(Y), n, 7).sum(2), index=Y.index, columns=ends)


def sf_weekly(W, cut):
    df = W.loc[:, :cut].stack().rename('y').reset_index(); df.columns = ['unique_id', 'ds', 'y']
    models = [AutoETS(season_length=1, alias='ETS_w'), AutoTheta(season_length=1, alias='Theta_w'),
              SimpleExponentialSmoothingOptimized(alias='SES_w'), ADIDA(alias='ADIDA_w'), IMAPA(alias='IMAPA_w'),
              TSB(alpha_d=0.2, alpha_p=0.2, alias='TSB_w'), CrostonSBA(alias='SBA_w')]
    fc = StatsForecast(models=models, freq='W-TUE', n_jobs=2).forecast(df=df, h=4)
    if 'unique_id' not in fc.columns: fc = fc.reset_index()
    return {m.alias: fc.pivot(index='unique_id', columns='ds', values=m.alias).reindex(W.index).values.clip(min=0) for m in models}


# ------------------------------------------------------------------ LightGBM, daily direct model
LGB = dict(objective='l1', n_estimators=500, learning_rate=0.04, num_leaves=31, min_child_samples=40, subsample=0.8,
           subsample_freq=1, colsample_bytree=0.8, reg_lambda=1.0, verbose=-1, n_jobs=2)


def _daily_feats(a, o, grp, sid):
    n = a.shape[0]; rows = []
    m7, m28, m56 = a[:, o - 6:o + 1].mean(1), a[:, o - 27:o + 1].mean(1), a[:, o - 55:o + 1].mean(1)
    scale = np.where(m28 > 0, m28, np.where(m56 > 0, m56, 1.0))
    z28 = (a[:, o - 27:o + 1] <= 0).mean(1)
    for h in range(1, H + 1):
        t = o + h; m = int(np.ceil(h / 7)); d = DATES[0] + pd.Timedelta(days=t)
        lags = np.stack([a[:, t - 7 * (m + j)] for j in range(4)], 1) / scale[:, None]
        wdm = np.stack([a[:, t - 7 * (m + j)] for j in range(8)], 1).mean(1) / np.where(m56 > 0, m56, 1.0)
        f = np.column_stack([np.full(n, h), np.full(n, d.dayofweek), np.full(n, d.day), np.full(n, int(d in HOLI)), lags,
                             m7 / scale, m56 / scale, z28, wdm, np.log1p(scale), grp, sid])
        rows.append(f)
    return np.stack(rows, 1), scale        # n x H x F


def lgb_daily(Y, cut, grp):
    a = Y.values.astype(float); n = len(Y); c = DATES.get_loc(cut)
    grp = np.zeros(n) if grp is None else pd.factorize(grp)[0]
    sid = np.arange(n) if n <= 20 else np.zeros(n)
    X, y, w = [], [], []
    for o in range(56, c - H + 1, 7):
        F, sc = _daily_feats(a, o, grp, sid)
        X.append(F.reshape(-1, F.shape[2])); y.append((a[:, o + 1:o + 1 + H] / sc[:, None]).reshape(-1)); w.append(np.repeat(sc, H))
    X, y, w = np.vstack(X), np.concatenate(y), np.concatenate(w)
    keep = w > 0
    mdl = lgb.LGBMRegressor(**LGB).fit(X[keep], y[keep], sample_weight=w[keep], categorical_feature=[13, 14] if n <= 20 else [13])
    F, sc = _daily_feats(a, c, grp, sid)
    p = mdl.predict(F.reshape(-1, F.shape[2])).reshape(n, H).clip(min=0) * sc[:, None]
    return p.reshape(n, 4, 7).sum(2)


# ------------------------------------------------------------------ LightGBM, weekly direct model
def _weekly_feats(a, o, ends, grp, sid):
    n = a.shape[0]; out = []
    m4 = a[:, o - 3:o + 1].mean(1); m13 = a[:, max(0, o - 12):o + 1].mean(1)
    scale = np.where(m4 > 0, m4, np.where(m13 > 0, m13, 1.0))
    lags = np.stack([a[:, o - j] for j in range(8)], 1) / scale[:, None]
    z8 = (a[:, o - 7:o + 1] <= 0).mean(1)
    for h in range(1, 5):
        e = ends[o] + pd.Timedelta(weeks=h); days = pd.date_range(e - pd.Timedelta(days=6), e)
        hol = int(any(d in HOLI for d in days)); pay = int(any(d.day == 27 for d in days))
        out.append(np.column_stack([np.full(n, h), np.full(n, hol), np.full(n, pay), np.full(n, e.month), lags, m13 / scale, z8,
                                    np.log1p(scale), grp, sid]))
    return np.stack(out, 1), scale


def lgb_weekly(W, cut, grp):
    a = W.values.astype(float); n = len(W); ends = W.columns; c = list(ends).index(cut)
    grp = np.zeros(n) if grp is None else pd.factorize(grp)[0]
    sid = np.arange(n) if n <= 20 else np.zeros(n)
    X, y, w = [], [], []
    for o in range(8, c - 4 + 1):
        F, sc = _weekly_feats(a, o, ends, grp, sid)
        X.append(F.reshape(-1, F.shape[2])); y.append((a[:, o + 1:o + 5] / sc[:, None]).reshape(-1)); w.append(np.repeat(sc, 4))
    X, y, w = np.vstack(X), np.concatenate(y), np.concatenate(w)
    keep = w > 0
    p = dict(LGB); p.update(n_estimators=300, min_child_samples=20)
    mdl = lgb.LGBMRegressor(**p).fit(X[keep], y[keep], sample_weight=w[keep], categorical_feature=[15, 16] if n <= 20 else [15])
    F, sc = _weekly_feats(a, c, ends, grp, sid)
    return mdl.predict(F.reshape(-1, F.shape[2])).reshape(n, 4).clip(min=0) * sc[:, None]


# ------------------------------------------------------------------ run every candidate at every origin
if __name__ == '__main__':
    FC, ACT = {}, {}
    for lv, Y in LEVELS.items():
        t0 = time.time(); W = to_weekly(Y); FC[lv] = {}; ACT[lv] = {}
        for cut in CUTS:
            ACT[lv][cut] = weekly_actual(Y, cut)
            fc = baselines(Y, cut)
            fc.update(sf_daily(Y, cut))
            fc.update(sf_weekly(W, cut))
            fc['LGBM_d'] = lgb_daily(Y, cut, GROUP[lv])
            fc['LGBM_w'] = lgb_weekly(W, cut, GROUP[lv])
            for m, M in fc.items():
                FC[lv].setdefault(m, {})[cut] = np.asarray(M, float)
        print(lv, len(Y), 'series', f'{time.time() - t0:.0f}s', flush=True)
        pickle.dump({'FC': FC, 'ACT': ACT, 'index': {k: list(v.index) for k, v in LEVELS.items()}}, open(f'{OUT}/fc_race.pkl', 'wb'))
    print('done')