"""Run the forecast race for one level (sku | van | principal) and save raw backtest results to WORK."""
import sys, time, pickle, numpy as np, pandas as pd
from pcommon import *
import fc_engine as E

level = sys.argv[1]
S = pd.read_parquet(f'{WORK}/sales.parquet'); P = pd.read_parquet(f'{WORK}/products.parquet')
S = S[(S.week_end >= FIRST_FULL_WEEK_END) & (S.week_end <= DATA_END)]
pr = P.set_index('sku')
if level == 'sku':
    S['uid'] = S.sku; val = 'qty_cases'
    cats = pr.l1_code.astype(str).to_dict()
elif level == 'van':
    S['uid'] = S.van; val = 'net_amt'; cats = None
elif level == 'principal':
    S['uid'] = S.sku.map(pr.l1_desc).fillna('Unknown principal'); val = 'net_amt'; cats = None
df = S.groupby(['uid', 'week_end'])[val].sum().reset_index().rename(columns={'uid': 'unique_id', 'week_end': 'ds', val: 'y'})
# full weekly grid; keep series with any activity in the training period
trained = df[df.ds <= TRAIN_END].groupby('unique_id').y.sum()
keep = trained[trained > 0].index
grid = pd.MultiIndex.from_product([keep, E.WEEKS], names=['unique_id', 'ds']).to_frame(index=False)
df = grid.merge(df, how='left').fillna({'y': 0.0})
df['y'] = df.y.clip(lower=0)
print(level, 'series', df.unique_id.nunique(), 'weeks', df.ds.nunique())
t0 = time.time(); cv1, cv12, fwd = E.stats_cv(df); print('stats done', round(time.time() - t0), 's')
t0 = time.time(); l1, l12, lf, imp = E.lgb_backtests(df, cats); print('lgb done', round(time.time() - t0), 's')
pickle.dump({'df': df, 'cv1': cv1, 'cv12': cv12, 'fwd': fwd, 'l1': l1, 'l12': l12, 'lf': lf, 'imp': imp}, open(f'{WORK}/race_{level}.pkl', 'wb'))
print(imp.head(10).to_string())