#!/usr/bin/env python3
# =============================================================================
# Author: Simon-Pierre Boucher
# Contact: contact@spboucher.ai
# =============================================================================
"""Step 03 — RQ2: IV term structure + smile vs GARCH/HAR-RV for RV forecasting.

In-sample panel R² for five nested models, per-ticker rolling out-of-sample
evaluation, and Diebold-Mariano tests of HAR-RV vs HAR-RV + IV surface.

Inputs : data/processed/merged_options_rv.parquet
Outputs: results/rq2_insample.csv, results/rq2_oos_results.csv,
         results/rq2_diebold_mariano.csv
"""

import warnings

import numpy as np
import pandas as pd

import _bootstrap  # noqa: F401
from wp7 import config
from wp7.data_io import load_merged
from wp7.econometrics import add_constant, ols, standardize, winsorize
from wp7.forecasting import diebold_mariano, ols_forecast, rolling_evaluation

warnings.filterwarnings('ignore')

TARGETS = {'1-Day RV': 'rv_fwd_1d', '5-Day RV': 'rv_fwd_5d'}

WINSORIZE_COLS = ['rv_fwd_1d', 'rv_fwd_5d', 'rv_lag1', 'rv_w', 'rv_m', 'sq_return',
                  'iv_atm_30d', 'iv_atm_90d', 'iv_term_slope', 'iv_skew_25d',
                  'implied_skewness', 'implied_kurtosis_proxy']


def main():
    print("=" * 70)
    print("RQ2: IV SURFACE vs GARCH/HAR-RV FOR RV FORECASTING")
    print("=" * 70)
    config.ensure_output_dirs()

    df = load_merged()
    df = df.sort_values(['ticker', 'trade_date']).reset_index(drop=True)
    print(f"Loaded {len(df):,} rows")

    # GARCH(1,1) proxy regressor: yesterday's squared daily return
    df['sq_return'] = df['daily_return'] ** 2

    # Per-ticker winsorization of every model input and target
    for col in WINSORIZE_COLS:
        if col in df.columns:
            df[col] = df.groupby('ticker')[col].transform(lambda x: winsorize(x))

    models = config.RQ2_MODELS

    # ── In-sample panel R² ──
    print("\n--- IN-SAMPLE PANEL REGRESSIONS ---")
    insample_results = []
    for target_label, target_col in TARGETS.items():
        for model_name, features in models.items():
            sub = df[features + [target_col, 'ticker']].dropna()
            if len(sub) < 100:
                continue

            X = add_constant(standardize(sub[features].values))
            y = sub[target_col].values
            coefs = ols(X, y)
            y_pred = X @ coefs
            ss_res = np.sum((y - y_pred) ** 2)
            ss_tot = np.sum((y - y.mean()) ** 2)
            r2 = 1 - ss_res / ss_tot
            adj_r2 = 1 - (1 - r2) * (len(y) - 1) / (len(y) - len(features) - 1)

            insample_results.append({
                'target': target_label,
                'model': model_name,
                'r2': r2,
                'adj_r2': adj_r2,
                'n_obs': len(sub),
                'n_features': len(features),
            })
            print(f"  {target_label} | {model_name}: R²={r2:.6f}, "
                  f"Adj-R²={adj_r2:.6f}, N={len(sub):,}")

    insample_df = pd.DataFrame(insample_results)
    insample_df.to_csv(config.RESULTS_DIR / "rq2_insample.csv", index=False)

    # ── Out-of-sample rolling evaluation (30 tickers with most data) ──
    print("\n--- OUT-OF-SAMPLE ROLLING EVALUATION ---")
    top_tickers = df.groupby('ticker').size().nlargest(30).index.tolist()
    df_top = df[df['ticker'].isin(top_tickers)]

    oos_results_all = []
    for target_label, target_col in TARGETS.items():
        print(f"\n  {target_label}:")
        oos = rolling_evaluation(df_top, models, target_col, target_label,
                                 window=500, step=250)
        oos_results_all.append(oos)
        if len(oos) > 0:
            summary = oos.groupby('model').agg({
                'avg_mse': 'mean', 'avg_mae': 'mean',
                'avg_r2_oos': 'mean', 'avg_qlike': 'mean',
                'n_windows': 'sum',
            }).round(6)
            print(summary.to_string())

    oos_df = pd.concat(oos_results_all, ignore_index=True)
    oos_df.to_csv(config.RESULTS_DIR / "rq2_oos_results.csv", index=False)

    # ── Diebold-Mariano: HAR-RV vs HAR-RV + IV surface (10 largest tickers) ──
    print("\n--- MODEL COMPARISON: DIEBOLD-MARIANO TESTS ---")
    f_har = ['rv_lag1', 'rv_w', 'rv_m']
    f_full = f_har + ['iv_atm_30d', 'iv_atm_90d', 'iv_term_slope',
                      'iv_skew_25d', 'implied_skewness', 'implied_kurtosis_proxy']
    dm_results = []
    for target_label, target_col in TARGETS.items():
        for ticker in top_tickers[:10]:
            td = df[df['ticker'] == ticker][f_full + [target_col]].dropna()
            if len(td) < 600:
                continue

            train, test = td.iloc[:500], td.iloc[500:]
            r1 = ols_forecast(train, test, f_har, target_col)
            r2 = ols_forecast(train, test, f_full, target_col)

            # Forecast errors on the test window (raw, before truncation at 0)
            X1 = np.column_stack([np.ones(len(test)), test[f_har].values])
            e1 = test[target_col].values - X1 @ r1['coefs']
            X2 = np.column_stack([np.ones(len(test)), test[f_full].values])
            e2 = test[target_col].values - X2 @ r2['coefs']

            dm_stat = diebold_mariano(e1, e2)
            dm_results.append({
                'ticker': ticker,
                'target': target_label,
                'mse_har': r1['mse'],
                'mse_har_iv': r2['mse'],
                'mse_improvement_pct': (r1['mse'] - r2['mse']) / r1['mse'] * 100,
                'dm_statistic': dm_stat,
                'dm_significant_5pct': abs(dm_stat) > 1.96,
            })

    dm_df = pd.DataFrame(dm_results)
    dm_df.to_csv(config.RESULTS_DIR / "rq2_diebold_mariano.csv", index=False)
    print(dm_df.to_string(index=False))

    print("\n" + "=" * 70)
    print("RQ2 SUMMARY")
    print("=" * 70)
    print("\nIn-sample R² comparison:")
    print(insample_df.pivot(index='model', columns='target', values='r2')
          .round(6).to_string())
    print("\nRQ2 COMPLETE.")


if __name__ == "__main__":
    main()