#!/usr/bin/env python3
# =============================================================================
# Author: Simon-Pierre Boucher
# Contact: contact@spboucher.ai
# =============================================================================
"""Step 08 — Subperiod and regime analysis: stability across market
conditions.

Sections A (nine subperiods), C (252-day rolling R²) and D (pre/post COVID)
run from the processed panel. Section B (VIX-regime conditioning) needs the
raw VIX series and is skipped gracefully when the raw stores are absent.

Inputs : data/processed/merged_options_rv.parquet
         (+ index_5min.duckdb for section B)
Outputs: results/subperiod_results.csv, results/regime_results.csv,
         results/rolling_r2.csv, results/pre_post_covid.csv
"""

import warnings

import pandas as pd

import _bootstrap  # noqa: F401
from wp7 import config
from wp7.data_io import RawDataUnavailableError, load_merged, load_vix_daily
from wp7.econometrics import pooled_regression_summary

warnings.filterwarnings('ignore')

FEATURES = config.RQ1_FEATURES
RV_MODELS = [('HAR-RV', config.HAR_FEATURES),
             ('HAR+IV', config.HAR_FEATURES + config.IV_FEATURES)]


def main():
    print("=" * 70)
    print("SUBPERIOD & REGIME ANALYSIS")
    print("=" * 70)
    config.ensure_output_dirs()

    df = load_merged()

    # ── A. Subperiod analysis ──
    print("\n--- A. SUBPERIOD ANALYSIS ---")
    subperiod_results = []
    for name, (start, end) in config.SUBPERIODS.items():
        sub_data = df[(df['trade_date'] >= start) & (df['trade_date'] <= end)]
        for target, horizon in [('ret_1d', '1D'), ('ret_5d', '5D')]:
            res = pooled_regression_summary(sub_data, FEATURES, target, f"{name}|{horizon}")
            if res:
                res['subperiod'] = name
                res['horizon'] = horizon
                subperiod_results.append(res)
                print(f"  {name} | {horizon}: R²={res['r2']:.6f}, "
                      f"N={res['n_obs']:,}, Sig={res['n_significant']}")

    print("\n  RV Forecasting by subperiod:")
    for name, (start, end) in config.SUBPERIODS.items():
        sub_data = df[(df['trade_date'] >= start) & (df['trade_date'] <= end)]
        for model_name, feats in RV_MODELS:
            res = pooled_regression_summary(sub_data, feats, 'rv_fwd_1d',
                                            f"{name}|{model_name}")
            if res:
                res['subperiod'] = name
                res['model'] = model_name
                subperiod_results.append(res)
                print(f"  {name} | {model_name}: R²={res['r2']:.6f}")

    pd.DataFrame(subperiod_results).to_csv(
        config.RESULTS_DIR / "subperiod_results.csv", index=False)

    # ── B. VIX regime analysis (needs raw VIX series) ──
    print("\n--- B. VIX REGIME ANALYSIS ---")
    try:
        vix = load_vix_daily()
        dfv = df.merge(vix, on='trade_date', how='left')
        dfv['vix_regime'] = pd.cut(dfv['vix_close'], bins=config.VIX_REGIME_BINS,
                                   labels=config.VIX_REGIME_LABELS)
        regime_results = []
        for regime in config.VIX_REGIME_LABELS:
            sub_data = dfv[dfv['vix_regime'] == regime]
            if len(sub_data) < 200:
                continue
            for target, horizon in [('ret_1d', '1D'), ('ret_5d', '5D'),
                                    ('rv_fwd_1d', 'RV_1D')]:
                feats = FEATURES if 'ret' in target else \
                    config.HAR_FEATURES + config.IV_FEATURES
                res = pooled_regression_summary(sub_data, feats, target,
                                                f"VIX_{regime}|{horizon}")
                if res:
                    res['regime'] = regime
                    res['target'] = horizon
                    regime_results.append(res)
                    print(f"  VIX {regime} | {horizon}: R²={res['r2']:.6f}, "
                          f"N={res['n_obs']:,}")
        if regime_results:
            pd.DataFrame(regime_results).to_csv(
                config.RESULTS_DIR / "regime_results.csv", index=False)
    except RawDataUnavailableError as exc:
        print(f"  [Section B skipped — raw stores unavailable; "
              f"existing regime_results.csv left untouched]\n  {exc}")

    # ── C. Rolling-window R² (252-day window, 63-day step) ──
    print("\n--- C. ROLLING WINDOW R² (252-day) ---")
    rolling_r2 = []
    dates_sorted = sorted(df['trade_date'].unique())
    for i in range(252, len(dates_sorted), 63):
        window_start = dates_sorted[max(0, i - 252)]
        window_end = dates_sorted[i]
        sub_data = df[(df['trade_date'] > window_start) & (df['trade_date'] <= window_end)]

        for target, horizon in [('ret_5d', '5D_Return'), ('rv_fwd_1d', '1D_RV')]:
            feats = FEATURES if target == 'ret_5d' else \
                config.HAR_FEATURES + config.IV_FEATURES
            res = pooled_regression_summary(sub_data, feats, target, "Rolling")
            if res:
                rolling_r2.append({'date': window_end, 'target': horizon,
                                   'r2': res['r2'], 'n_obs': res['n_obs']})

    rolling_df = pd.DataFrame(rolling_r2)
    rolling_df.to_csv(config.RESULTS_DIR / "rolling_r2.csv", index=False)
    print(f"  {len(rolling_df)} rolling windows computed")
    for target in rolling_df['target'].unique():
        sub = rolling_df[rolling_df['target'] == target]
        print(f"\n  {target}: mean R²={sub['r2'].mean():.6f}, "
              f"min={sub['r2'].min():.6f}, max={sub['r2'].max():.6f}, "
              f"std={sub['r2'].std():.6f}")

    # ── D. Pre vs post COVID ──
    print("\n--- D. PRE vs POST COVID COMPARISON ---")
    pre_covid = df[df['trade_date'] < '2020-03-01']
    post_covid = df[df['trade_date'] >= '2020-03-01']

    comparison_rows = []
    for period_name, period_data in [('Pre-COVID', pre_covid), ('Post-COVID', post_covid)]:
        for target, horizon in [('ret_1d', '1D'), ('ret_5d', '5D')]:
            res = pooled_regression_summary(period_data, FEATURES, target,
                                            f"{period_name}|{horizon}")
            if res:
                comparison_rows.append({**res, 'period': period_name, 'target': horizon})
                print(f"  {period_name} | {horizon}: R²={res['r2']:.6f}")
        for model_name, feats in RV_MODELS:
            res = pooled_regression_summary(period_data, feats, 'rv_fwd_1d',
                                            f"{period_name}|{model_name}")
            if res:
                comparison_rows.append({**res, 'period': period_name, 'target': model_name})
                print(f"  {period_name} | {model_name}: R²={res['r2']:.6f}")

    pd.DataFrame(comparison_rows).to_csv(
        config.RESULTS_DIR / "pre_post_covid.csv", index=False)

    print("\nSUBPERIOD & REGIME ANALYSIS COMPLETE.")


if __name__ == "__main__":
    main()