"""
Polleo Walk-Forward Backtest — v3.6 (channel-split)
====================================================
Generates 1-step-ahead forecasts for the last N weeks using the same
channel-split logic as forecast_engine.py v3.6:

- For each SKU, runs 3 time series through StatsForecast: total, retail, wholesale
- Adds seasonal models (season_length=4) on the wholesale channel to catch
  regular monthly buyer patterns
- For wholesale-dominant SKUs (ws_share >= 0.5): forecast = best_retail +
  best_wholesale (channel-split mode)
- For retail-dominant SKUs: forecast = best_total (current mode)

Results saved to data/backtest_fa.csv with additional columns:
  forecast_retail, forecast_wholesale, actual_retail, actual_wholesale,
  channel_mode, ws_share, model_retail, model_wholesale
Existing columns (sku, year, week, forecast, actual, model, cat, oznaka)
are preserved for backward compatibility with the FA dashboard.

Usage: python run_backtest.py [n_weeks] [data_dir]
       Default: 8 weeks, data/ directory
"""

import pandas as pd, numpy as np, sys, os
from datetime import datetime

# Fix Windows console encoding
import io
if sys.stdout.encoding and sys.stdout.encoding.lower() not in ('utf-8', 'utf8'):
    try:
        sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
        sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')
    except Exception:
        pass

try:
    from statsforecast import StatsForecast as SF
    from statsforecast.models import (
        AutoARIMA, AutoCES, AutoTheta, Naive,
        CrostonOptimized, ADIDA, IMAPA, TSB
    )
    HAS_SF = True
except ImportError:
    HAS_SF = False


def run_backtest(n_weeks=8, data_dir='data'):
    """Run channel-split walk-forward backtest for the last n_weeks completed weeks."""
    print(f'\n{"="*60}')
    print(f'  POLLEO WALK-FORWARD BACKTEST (v3.6 channel-split)')
    print(f'  {datetime.now().strftime("%Y-%m-%d %H:%M")}')
    print(f'{"="*60}')

    if not HAS_SF:
        print('\n  ERROR: statsforecast not installed. pip install statsforecast')
        return None

    sales_path = os.path.join(data_dir, 'sales_clean.csv')
    plan_path = os.path.join(data_dir, 'sku_plan_list.csv')

    if not os.path.exists(sales_path):
        print(f'  ERROR: {sales_path} not found')
        return None
    if not os.path.exists(plan_path):
        print(f'  ERROR: {plan_path} not found')
        return None

    sc = pd.read_csv(sales_path)
    plan = pd.read_csv(plan_path)
    sku_list = plan['sku'].tolist()
    sku_meta = {}
    for _, r in plan.iterrows():
        sku_meta[r['sku']] = {
            'cat': r.get('cat', ''),
            'ozn': r.get('oznaka', ''),
            'name': r.get('name', ''),
        }

    iso = datetime.now().isocalendar()
    cy, cw = iso[0], iso[1]
    current_key = cy * 100 + cw

    all_yw = sc[['year', 'week']].drop_duplicates().sort_values(['year', 'week'])
    all_yw['yw_key'] = all_yw['year'] * 100 + all_yw['week']
    all_yw = all_yw[all_yw['yw_key'] < current_key]
    all_weeks = list(zip(all_yw['year'].astype(int), all_yw['week'].astype(int)))

    if len(all_weeks) < n_weeks + 4:
        print(f'  ERROR: Not enough history. Need {n_weeks + 4} weeks, have {len(all_weeks)}')
        return None

    target_weeks = all_weeks[-n_weeks:]
    print(f'  History: {len(all_weeks)} weeks')
    print(f'  Backtest targets: CW{target_weeks[0][1]} ({target_weeks[0][0]}) to CW{target_weeks[-1][1]} ({target_weeks[-1][0]})')
    print(f'  SKUs: {len(sku_list)}')

    # ---- Build channel-split time series per SKU ----
    sc_indexed = sc.set_index(['sku', 'year', 'week'])
    sf_rows = []
    sku_series = {}   # sku -> {'total': arr, 'retail': arr, 'wholesale': arr}

    ws_cleaned = 0
    # Spikes inside the test window are left alone so FA reflects real error;
    # only pre-window history is cleaned for training.
    train_cutoff_idx = len(all_weeks) - n_weeks
    for sku in sku_list:
        vals_t, vals_r, vals_w, vals_sp = [], [], [], []
        for y, w in all_weeks:
            try:
                row = sc_indexed.loc[(sku, y, w)]
                if isinstance(row, pd.DataFrame):
                    row = row.iloc[0]
                qt = float(row['qty_total'])
                qr = float(row.get('qty_retail', 0) or 0) + float(row.get('qty_webshop', 0) or 0)
                qw = float(row.get('qty_wholesale', 0) or 0)
                sp = int(row.get('is_wholesale_spike', 0) or 0)
            except KeyError:
                qt, qr, qw, sp = 0.0, 0.0, 0.0, 0
            vals_t.append(qt)
            vals_r.append(qr)
            vals_w.append(qw)
            vals_sp.append(sp)

        if sum(1 for v in vals_t if v > 0) < 4:
            continue

        # v3.7: Clean wholesale spikes out of training history so the baseline
        # model doesn't learn phantom recurring surges from one-shot sell-in.
        if any(vals_sp):
            wh_arr = np.array(vals_w, dtype=float)
            sp_arr = np.array(vals_sp)
            nonspike_wh = wh_arr[(sp_arr == 0) & (wh_arr > 0)]
            median_wh = float(np.median(nonspike_wh)) if len(nonspike_wh) >= 4 else 0.0
            for t, f in enumerate(vals_sp):
                if f and t < train_cutoff_idx:
                    retail_part = max(vals_t[t] - vals_w[t], 0.0)
                    vals_w[t] = median_wh
                    vals_t[t] = retail_part + median_wh
                    ws_cleaned += 1

        sku_series[sku] = {
            'total': np.array(vals_t),
            'retail': np.array(vals_r),
            'wholesale': np.array(vals_w),
        }
        for t, v in enumerate(vals_t):
            sf_rows.append({'unique_id': sku, 'ds': t, 'y': max(v, 0.0)})
        for t, v in enumerate(vals_r):
            sf_rows.append({'unique_id': f'{sku}_R', 'ds': t, 'y': max(v, 0.0)})
        for t, v in enumerate(vals_w):
            sf_rows.append({'unique_id': f'{sku}_W', 'ds': t, 'y': max(v, 0.0)})
    print(f'  Wholesale spike weeks cleaned (training-only): {ws_cleaned}')

    sf_df = pd.DataFrame(sf_rows)
    n_series = sf_df['unique_id'].nunique()
    print(f'\n  SF DataFrame: {len(sf_df):,} rows, {n_series} series (total + retail + wholesale)')

    # ---- Main CV: fast models on all 3 channels ----
    models = [
        AutoCES(season_length=1),
        AutoTheta(season_length=1),
        CrostonOptimized(),
        ADIDA(),
        IMAPA(),
        TSB(alpha_d=0.2, alpha_p=0.2),
    ]
    sf = SF(models=models, freq=1, n_jobs=1, fallback_model=Naive())
    print(f'  Running cross_validation (h=1, step_size=1, n_windows={n_weeks})...')

    import time
    t0 = time.time()
    try:
        cv = sf.cross_validation(df=sf_df, h=1, step_size=1, n_windows=n_weeks)
    except Exception as e:
        print(f'  ERROR in cross_validation: {e}')
        import traceback; traceback.print_exc()
        return None
    print(f'  Main CV done in {time.time() - t0:.1f}s')

    # ---- Seasonal wholesale CV (season_length=4, monthly pattern) ----
    cv_seasonal = None
    seasonal_cols = []
    sf_df_w = sf_df[sf_df['unique_id'].str.endswith('_W')].copy()
    if len(sf_df_w) > 0:
        print(f'  Running seasonal wholesale CV (season_length=4)...')
        models_seasonal = [
            AutoARIMA(season_length=4),
            AutoTheta(season_length=4),
        ]
        sf_seasonal = SF(models=models_seasonal, freq=1, n_jobs=1, fallback_model=Naive())
        t1 = time.time()
        try:
            cv_seasonal = sf_seasonal.cross_validation(
                df=sf_df_w, h=1, step_size=1, n_windows=n_weeks
            )
            raw_seasonal_cols = [
                c for c in cv_seasonal.columns
                if c not in ('unique_id', 'ds', 'y', 'cutoff')
            ]
            rename_map = {c: f'{c}_s4' for c in raw_seasonal_cols}
            cv_seasonal = cv_seasonal.rename(columns=rename_map)
            seasonal_cols = list(rename_map.values())
            print(f'  Seasonal CV done in {time.time() - t1:.1f}s')
        except Exception as e:
            print(f'  Seasonal CV failed ({type(e).__name__}: {e}) — continuing without seasonal')
            cv_seasonal = None

    model_cols = [c for c in cv.columns if c not in ('unique_id', 'ds', 'y', 'cutoff')]
    print(f'  Models: {model_cols}' + (f' + seasonal-ws: {seasonal_cols}' if seasonal_cols else ''))

    # ---- Build results per SKU per target week ----
    results = []
    n_total = len(all_weeks)

    for sku in sku_series:
        meta = sku_meta.get(sku, {'cat': '', 'ozn': '', 'name': ''})
        total_series = sku_series[sku]['total']
        retail_series = sku_series[sku]['retail']
        wholesale_series = sku_series[sku]['wholesale']

        tot_hist_sum = retail_series.sum() + wholesale_series.sum()
        ws_share = float(wholesale_series.sum() / tot_hist_sum) if tot_hist_sum > 0 else 0.0

        cv_t = cv[cv['unique_id'] == sku].sort_values('ds').reset_index(drop=True)
        cv_r = cv[cv['unique_id'] == f'{sku}_R'].sort_values('ds').reset_index(drop=True)
        cv_w = cv[cv['unique_id'] == f'{sku}_W'].sort_values('ds').reset_index(drop=True)

        if len(cv_t) == 0:
            continue

        # Merge seasonal wholesale into cv_w
        if cv_seasonal is not None:
            cv_w_s = cv_seasonal[cv_seasonal['unique_id'] == f'{sku}_W'] \
                        .sort_values('ds').reset_index(drop=True)
            if len(cv_w_s) > 0 and len(cv_w) > 0:
                cv_w = cv_w.merge(
                    cv_w_s[['ds'] + seasonal_cols],
                    on='ds', how='left'
                )

        actuals_t = cv_t['y'].values
        actuals_r_arr = retail_series[-n_weeks:]
        actuals_w_arr = wholesale_series[-n_weeks:]

        def _pick_best(cv_df, model_names, actuals, penalty_suffix=None, penalty=1.10):
            if cv_df is None or len(cv_df) == 0:
                return None, None
            valid = actuals > 0
            if valid.sum() == 0:
                return None, None
            scores = {}
            for mc in model_names:
                if mc not in cv_df.columns:
                    continue
                col = cv_df[mc].values
                preds = np.maximum(np.nan_to_num(col, nan=0.0), 0)
                if len(preds) != len(actuals):
                    continue
                pen = penalty if (penalty_suffix and mc.endswith(penalty_suffix)) else 1.0
                scores[mc] = np.mean(np.abs(preds[valid] - actuals[valid])) * pen
            if not scores:
                return None, None
            best = min(scores, key=scores.get)
            return best, np.maximum(np.nan_to_num(cv_df[best].values, nan=0.0), 0)

        best_t, preds_t = _pick_best(cv_t, model_cols, actuals_t)
        if best_t is None:
            continue

        best_r, preds_r = _pick_best(cv_r, model_cols, actuals_r_arr)
        ws_model_names = list(model_cols) + seasonal_cols
        best_w, preds_w = _pick_best(
            cv_w, ws_model_names, actuals_w_arr,
            penalty_suffix='_s4', penalty=1.10
        )

        use_split = (
            ws_share >= 0.5
            and preds_r is not None and preds_w is not None
            and len(preds_r) == len(actuals_t) and len(preds_w) == len(actuals_t)
        )
        channel_mode = 'split' if use_split else 'total'

        for i in range(len(target_weeks)):
            yr, wk = target_weeks[i]
            actual_t = round(actuals_t[i]) if i < len(actuals_t) else 0
            actual_r = round(actuals_r_arr[i]) if i < len(actuals_r_arr) else 0
            actual_w = round(actuals_w_arr[i]) if i < len(actuals_w_arr) else 0

            train_end = n_total - n_weeks + i
            t_hist = total_series[max(0, train_end - 13):train_end]
            r_hist = retail_series[max(0, train_end - 13):train_end]
            w_hist = wholesale_series[max(0, train_end - 26):train_end]

            def _cap_floor(fc, hist, for_wholesale=False):
                nz = hist[hist > 0]
                if len(nz) >= 3:
                    if for_wholesale:
                        fc = min(fc, float(np.max(nz)) * 1.5)
                    else:
                        fc = min(fc, float(np.mean(nz)) * 2.0)
                        fc = max(fc, float(np.median(nz)) * 0.5)
                return max(0, fc)

            fc_r_raw = float(preds_r[i]) if preds_r is not None and i < len(preds_r) else 0.0
            fc_w_raw = float(preds_w[i]) if preds_w is not None and i < len(preds_w) else 0.0
            fc_r = _cap_floor(fc_r_raw, r_hist, for_wholesale=False)
            fc_w = _cap_floor(fc_w_raw, w_hist, for_wholesale=True)

            if use_split:
                fc_total = fc_r + fc_w
            else:
                fc_t_raw = float(preds_t[i])
                fc_total = _cap_floor(fc_t_raw, t_hist, for_wholesale=False)

            results.append({
                'sku': sku, 'year': yr, 'week': wk,
                'forecast': round(fc_total),
                'actual': actual_t,
                'forecast_retail': round(fc_r),
                'forecast_wholesale': round(fc_w),
                'actual_retail': actual_r,
                'actual_wholesale': actual_w,
                'channel_mode': channel_mode,
                'ws_share': round(ws_share, 3),
                'model': best_t if not use_split else f'{best_r}+{best_w}',
                'model_retail': best_r or '',
                'model_wholesale': best_w or '',
                'cat': meta['cat'], 'oznaka': meta['ozn'],
            })

    result_df = pd.DataFrame(results)
    out_path = os.path.join(data_dir, 'backtest_fa.csv')

    # Merge with existing backtest history: keep rows for other (sku, year, week)
    # keys, overwrite any that this run just recomputed. Prevents short runs
    # (e.g. n_weeks=1) from wiping the earlier accuracy record.
    if os.path.exists(out_path):
        try:
            prev = pd.read_csv(out_path)
            new_keys = set(zip(result_df['sku'], result_df['year'], result_df['week']))
            prev_keep = prev[~prev.apply(
                lambda r: (r['sku'], r['year'], r['week']) in new_keys, axis=1)]
            result_df = pd.concat([prev_keep, result_df], ignore_index=True)
            result_df = result_df.sort_values(['year', 'week', 'sku']).reset_index(drop=True)
        except Exception as e:
            print(f'  [merge] skipped (falling back to overwrite): {e}')
    result_df.to_csv(out_path, index=False)

    # ---- Summary ----
    print(f'\n  Results:')
    valid = result_df[result_df['actual'] > 0].copy()
    if len(valid) > 0:
        valid['error'] = np.abs(valid['forecast'] - valid['actual'])
        valid['fa'] = np.maximum(0, 1 - valid['error'] / valid['actual'])
        valid['hit'] = (valid['error'] / valid['actual'].clip(lower=1) <= 0.3).astype(int)

        for ozn in sorted(valid['oznaka'].unique()):
            sub = valid[valid['oznaka'] == ozn]
            fa = sub['fa'].mean() * 100
            bias = (sub['forecast'].sum() - sub['actual'].sum()) / max(sub['actual'].sum(), 1) * 100
            hr = sub['hit'].mean() * 100
            print(f'  {ozn}: FA={fa:.1f}%, Bias={bias:+.1f}%, HR={hr:.0f}%, {sub["sku"].nunique()} SKUs')

        total_fa = valid['fa'].mean() * 100
        total_bias = (valid['forecast'].sum() - valid['actual'].sum()) / max(valid['actual'].sum(), 1) * 100
        print(f'  TOTAL: FA={total_fa:.1f}%, Bias={total_bias:+.1f}%')

        print(f'\n  By channel mode:')
        for mode in sorted(valid['channel_mode'].unique()):
            sub = valid[valid['channel_mode'] == mode]
            fa = sub['fa'].mean() * 100
            bias = (sub['forecast'].sum() - sub['actual'].sum()) / max(sub['actual'].sum(), 1) * 100
            n_sku = sub['sku'].nunique()
            print(f'    {mode}: FA={fa:.1f}%, Bias={bias:+.1f}%, {n_sku} SKUs, {len(sub)} rows')

        split_rows = valid[valid['channel_mode'] == 'split'].copy()
        if len(split_rows) > 0:
            print(f'\n  Channel FA (split-mode SKUs only):')
            for ch_label, fc_col, act_col in [
                ('  RETAIL   ', 'forecast_retail', 'actual_retail'),
                ('  WHOLESALE', 'forecast_wholesale', 'actual_wholesale'),
            ]:
                rows = split_rows[split_rows[act_col] > 0].copy()
                if len(rows) == 0:
                    print(f'  {ch_label}: (no positive actuals)')
                    continue
                err = np.abs(rows[fc_col] - rows[act_col]).sum()
                act = rows[act_col].sum()
                fc = rows[fc_col].sum()
                fa = max(0, (1 - err / act) * 100) if act > 0 else 0
                bias = ((fc - act) / act * 100) if act > 0 else 0
                print(f'  {ch_label}: FA={fa:.1f}%, Bias={bias:+.1f}%, {rows["sku"].nunique()} SKUs')

    print(f'\n  Saved: {out_path} ({len(result_df):,} rows)')
    print(f'{"="*60}\n')
    return result_df


if __name__ == '__main__':
    n = int(sys.argv[1]) if len(sys.argv) > 1 else 8
    data_dir = sys.argv[2] if len(sys.argv) > 2 else 'data'
    try:
        run_backtest(n, data_dir)
    except Exception as e:
        print(f'\n  ERROR: {e}')
        import traceback; traceback.print_exc()
    try:
        input('\n  Press Enter to close...')
    except (EOFError, OSError):
        pass