#!/usr/bin/env python3
# =============================================================================
# Author: Simon-Pierre Boucher
# Contact: contact@spboucher.ai
# =============================================================================
"""Step 10 — Robustness checks.

A. Newey-West HAC(5) standard errors
B. Double-clustered (ticker + date) standard errors, CGM (2011)
C. Additional control variables (volume, OI, lagged returns)
D. Quantile regressions (IRLS) at τ ∈ {0.10, 0.25, 0.50, 0.75, 0.90}
E. Excluding high-VIX periods (needs the raw VIX series; console-only)
F. Ticker-by-ticker R² distribution

Note: the original implementation of section D materialized an n×n diagonal
weight matrix (~97 GB at this sample size) and could not complete; the IRLS
weighting is now applied by broadcasting — mathematically identical — so
sections D and F produce the two result files that were missing from the
original archive (see AUDIT.md §6.2).

Inputs : data/processed/merged_options_rv.parquet
         (+ index_5min.duckdb for section E)
Outputs: results/robustness_newey_west.csv,
         results/robustness_double_clustered.csv,
         results/robustness_with_controls.csv,
         results/robustness_quantile_regression.csv,
         results/robustness_ticker_r2.csv
"""

import warnings

import numpy as np
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 (add_constant, double_clustered_tstats, hc1_tstats,
                              newey_west_tstats, ols, quantile_regression,
                              r_squared, standardize, winsorize)

warnings.filterwarnings('ignore')

FEATURES = config.RQ1_FEATURES
HORIZONS = [('ret_1d', '1-Day'), ('ret_5d', '5-Day')]


def prepare(data: pd.DataFrame, features: list, target: str,
            keep: tuple = ()) -> pd.DataFrame:
    """Drop incomplete rows and winsorize features and target at 1%/99%."""
    sub = data[features + [target, *keep]].dropna()
    for f in features:
        sub[f] = winsorize(sub[f])
    sub[target] = winsorize(sub[target])
    return sub


def main():
    print("=" * 70)
    print("ROBUSTNESS CHECKS")
    print("=" * 70)
    config.ensure_output_dirs()

    df = load_merged()

    # ── A. Newey-West HAC standard errors ──
    print("\n--- A. NEWEY-WEST HAC STANDARD ERRORS (lag=5) ---")
    nw_results = []
    for target, horizon in HORIZONS:
        sub = prepare(df, FEATURES, target)
        X = add_constant(standardize(sub[FEATURES].values))
        y = sub[target].values
        coefs = ols(X, y)
        se, t_stats = newey_west_tstats(X, y, coefs, n_lags=5)
        r2 = r_squared(y, X @ coefs)

        res_df = pd.DataFrame({
            'variable': ['const'] + FEATURES,
            'coefficient': coefs,
            'nw_se': se,
            'nw_t_stat': t_stats,
            'nw_sig_5pct': np.abs(t_stats) > 1.96,
        })
        res_df['target'] = horizon
        res_df['method'] = 'Newey-West(5)'
        nw_results.append(res_df)

        sig = res_df[(res_df['nw_sig_5pct']) & (res_df['variable'] != 'const')]
        print(f"\n  {horizon}: R²={r2:.6f}, N={len(sub):,}")
        print(f"  Significant (NW): {', '.join(sig['variable'].tolist())}")
        for _, row in res_df.iterrows():
            star = "**" if abs(row['nw_t_stat']) > 2.576 else \
                ("*" if abs(row['nw_t_stat']) > 1.96 else "")
            print(f"    {row['variable']:25s}: β={row['coefficient']:9.6f}  "
                  f"t_NW={row['nw_t_stat']:7.3f} {star}")

    pd.concat(nw_results, ignore_index=True).to_csv(
        config.RESULTS_DIR / "robustness_newey_west.csv", index=False)

    # ── B. Double-clustered standard errors (ticker + date) ──
    print("\n--- B. DOUBLE-CLUSTERED SE (ticker + date) ---")
    dc_results = []
    for target, horizon in HORIZONS:
        sub = prepare(df, FEATURES, target, keep=('ticker', 'trade_date'))
        X = add_constant(standardize(sub[FEATURES].values))
        y = sub[target].values
        coefs = ols(X, y)
        se, t_stats = double_clustered_tstats(
            X, y, coefs, sub['ticker'].values, sub['trade_date'].values)
        r2 = r_squared(y, X @ coefs)

        res_df = pd.DataFrame({
            'variable': ['const'] + FEATURES,
            'coefficient': coefs,
            'dc_se': se,
            'dc_t_stat': t_stats,
            'dc_sig_5pct': np.abs(t_stats) > 1.96,
        })
        res_df['target'] = horizon
        dc_results.append(res_df)

        sig = res_df[(res_df['dc_sig_5pct']) & (res_df['variable'] != 'const')]
        print(f"\n  {horizon}: R²={r2:.6f}, N={len(sub):,}")
        print(f"  Significant (DC): {', '.join(sig['variable'].tolist())}")
        for _, row in res_df.iterrows():
            star = "**" if abs(row['dc_t_stat']) > 2.576 else \
                ("*" if abs(row['dc_t_stat']) > 1.96 else "")
            print(f"    {row['variable']:25s}: β={row['coefficient']:9.6f}  "
                  f"t_DC={row['dc_t_stat']:7.3f} {star}")

    pd.concat(dc_results, ignore_index=True).to_csv(
        config.RESULTS_DIR / "robustness_double_clustered.csv", index=False)

    # ── C. Additional control variables ──
    print("\n--- C. WITH CONTROL VARIABLES (volume, log_oi, spread proxy) ---")
    df['log_volume'] = np.log1p(df['total_option_volume'])
    df['log_oi'] = np.log1p(df['total_oi'])
    df['abs_return'] = df['daily_return'].abs()
    df['ret_lag1'] = df.groupby('ticker')['daily_return'].shift(1)
    df['ret_lag5'] = df.groupby('ticker')['daily_return'].transform(
        lambda x: x.shift(1).rolling(5).sum())

    controls = ['log_volume', 'log_oi', 'abs_return', 'ret_lag1', 'ret_lag5']
    features_ctrl = FEATURES + controls

    ctrl_results = []
    for target, horizon in HORIZONS:
        sub = prepare(df, features_ctrl, target)
        X = add_constant(standardize(sub[features_ctrl].values))
        y = sub[target].values
        coefs = ols(X, y)
        r2 = r_squared(y, X @ coefs)
        _, t = hc1_tstats(X, y, coefs, len(features_ctrl))

        print(f"\n  {horizon} with controls: R²={r2:.6f}, N={len(sub):,}")
        for i, name in enumerate(['const'] + features_ctrl):
            star = "**" if abs(t[i]) > 2.576 else ("*" if abs(t[i]) > 1.96 else "")
            print(f"    {name:25s}: β={coefs[i]:9.6f}  t={t[i]:7.3f} {star}")
            ctrl_results.append({'target': horizon, 'variable': name,
                                 'coefficient': coefs[i], 't_stat': t[i], 'r2': r2})

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

    # ── D. Quantile regressions ──
    print("\n--- D. QUANTILE REGRESSION (iterative reweighting) ---")
    qr_results = []
    for target, horizon in [('ret_5d', '5-Day')]:
        sub = prepare(df, FEATURES, target)
        X = add_constant(standardize(sub[FEATURES].values))
        y = sub[target].values

        for tau in [0.10, 0.25, 0.50, 0.75, 0.90]:
            coefs = quantile_regression(X, y, tau)
            print(f"\n  {horizon} | τ={tau}:")
            for i, name in enumerate(['const'] + FEATURES):
                qr_results.append({'target': horizon, 'tau': tau,
                                   'variable': name, 'coefficient': coefs[i]})
                print(f"    {name:25s}: β={coefs[i]:9.6f}")

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

    # ── E. Excluding high-VIX periods (console-only, needs raw VIX) ──
    print("\n--- E. EXCLUDING HIGH-VIX PERIODS (VIX < 30) ---")
    try:
        vix = load_vix_daily()
        calm = df.merge(vix, on='trade_date', how='left')
        calm = calm[calm['vix_close'] < 30]
        for target, horizon in HORIZONS:
            sub = prepare(calm, FEATURES, target)
            X = add_constant(standardize(sub[FEATURES].values))
            y = sub[target].values
            coefs = ols(X, y)
            print(f"  {horizon} (VIX<30): R²={r_squared(y, X @ coefs):.6f}, N={len(sub):,}")
    except RawDataUnavailableError as exc:
        print(f"  [Section E skipped — raw stores unavailable]\n  {exc}")

    # ── F. Ticker-by-ticker R² distribution (5-day returns) ──
    print("\n--- F. TICKER-BY-TICKER R² DISTRIBUTION (5-Day returns) ---")
    ticker_r2 = []
    for ticker in df['ticker'].unique():
        td = df[df['ticker'] == ticker]
        sub = td[FEATURES + ['ret_5d']].dropna()
        if len(sub) < 200:
            continue
        for f in FEATURES:
            sub[f] = winsorize(sub[f])
        sub['ret_5d'] = winsorize(sub['ret_5d'])
        X = add_constant(standardize(sub[FEATURES].values))
        y = sub['ret_5d'].values
        coefs = ols(X, y)
        ticker_r2.append({'ticker': ticker, 'r2_5d': r_squared(y, X @ coefs),
                          'n_obs': len(sub)})

    ticker_r2_df = pd.DataFrame(ticker_r2)
    ticker_r2_df.to_csv(config.RESULTS_DIR / "robustness_ticker_r2.csv", index=False)
    print(f"  Distribution of R² across {len(ticker_r2_df)} tickers:")
    for stat, val in [('Mean', ticker_r2_df['r2_5d'].mean()),
                      ('Median', ticker_r2_df['r2_5d'].median()),
                      ('Std', ticker_r2_df['r2_5d'].std()),
                      ('Min', ticker_r2_df['r2_5d'].min()),
                      ('Max', ticker_r2_df['r2_5d'].max()),
                      ('Q25', ticker_r2_df['r2_5d'].quantile(0.25)),
                      ('Q75', ticker_r2_df['r2_5d'].quantile(0.75))]:
        print(f"    {stat}:   {val:.6f}")

    print("\nROBUSTNESS CHECKS COMPLETE.")


if __name__ == "__main__":
    main()