#!/usr/bin/env python3
# =============================================================================
# Author: Simon-Pierre Boucher
# Contact: contact@spboucher.ai
# =============================================================================
"""Step 11 — Granger causality and VAR analysis.

Lead-lag structure between option-implied measures and returns / realized
volatility: per-ticker bivariate Granger F-tests (5 lags), bivariate VAR(5)
with impulse-response functions, and forecast-error variance decomposition.

Inputs : data/processed/merged_options_rv.parquet
Outputs: results/granger_causality.csv, results/var_results.csv,
         results/irf_results.csv, results/fevd_results.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 estimate_var, granger_test, winsorize

warnings.filterwarnings('ignore')

TEST_PAIRS = [
    ('iv_atm_30d', 'rv_daily', 'ATM_IV → RV'),
    ('rv_daily', 'iv_atm_30d', 'RV → ATM_IV'),
    ('iv_skew_25d', 'rv_daily', 'Skew → RV'),
    ('rv_daily', 'iv_skew_25d', 'RV → Skew'),
    ('iv_atm_30d', 'daily_return', 'ATM_IV → Return'),
    ('daily_return', 'iv_atm_30d', 'Return → ATM_IV'),
    ('iv_skew_25d', 'daily_return', 'Skew → Return'),
    ('daily_return', 'iv_skew_25d', 'Return → Skew'),
    ('pc_volume_ratio', 'daily_return', 'PC_Ratio → Return'),
    ('daily_return', 'pc_volume_ratio', 'Return → PC_Ratio'),
    ('implied_skewness', 'rv_daily', 'Impl_Skew → RV'),
    ('implied_kurtosis_proxy', 'rv_daily', 'Impl_Kurt → RV'),
]

IRF_PERIODS = 20


def main():
    print("=" * 70)
    print("GRANGER CAUSALITY & VAR ANALYSIS")
    print("=" * 70)
    config.ensure_output_dirs()

    df = load_merged()
    df = df.sort_values(['ticker', 'trade_date'])

    # ── A. Granger causality tests (per ticker, averaged) ──
    print("\n--- A. GRANGER CAUSALITY TESTS ---")
    granger_results = []
    for x_var, y_var, label in TEST_PAIRS:
        ticker_results = []
        for ticker in df['ticker'].unique():
            td = df[df['ticker'] == ticker][[x_var, y_var]].dropna()
            if len(td) < 100:
                continue
            x = winsorize(td[x_var]).values
            y = winsorize(td[y_var]).values
            res = granger_test(y, x, max_lag=5)
            if res:
                ticker_results.append(res)

        if ticker_results:
            avg_f = np.mean([r['f_stat'] for r in ticker_results])
            avg_p = np.mean([r['p_value'] for r in ticker_results])
            pct_sig = np.mean([1 if r['p_value'] < 0.05 else 0 for r in ticker_results])
            granger_results.append({
                'test': label,
                'x_causes_y': f"{x_var} → {y_var}",
                'avg_f_stat': avg_f,
                'avg_p_value': avg_p,
                'pct_significant_5pct': pct_sig,
                'n_tickers': len(ticker_results),
            })
            sig_str = "***" if avg_p < 0.01 else \
                ("**" if avg_p < 0.05 else ("*" if avg_p < 0.10 else ""))
            print(f"  {label:25s}: F={avg_f:8.3f}, p={avg_p:.4f}, "
                  f"{pct_sig*100:.1f}% sig  {sig_str}")

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

    # ── B. Bivariate VAR(5): ATM IV ↔ RV (20 largest tickers) ──
    print("\n--- B. BIVARIATE VAR: IV_ATM ↔ RV (pooled) ---")
    var_results, irf_all = [], []
    top_tickers = df.groupby('ticker').size().nlargest(20).index.tolist()

    for ticker in top_tickers:
        td = df[df['ticker'] == ticker][['iv_atm_30d', 'rv_daily']].dropna()
        if len(td) < 200:
            continue
        y1 = winsorize(td['iv_atm_30d']).values
        y2 = winsorize(td['rv_daily']).values
        y1 = (y1 - y1.mean()) / y1.std()
        y2 = (y2 - y2.mean()) / y2.std()

        res = estimate_var(y1, y2, lags=5, irf_periods=IRF_PERIODS)
        if res:
            var_results.append({
                'ticker': ticker,
                'r2_iv_eq': res['r2_eq1'],
                'r2_rv_eq': res['r2_eq2'],
                'n_obs': res['n_obs'],
            })
            for h in range(IRF_PERIODS):
                irf_all.append({
                    'ticker': ticker, 'horizon': h,
                    'iv_to_iv': res['irf'][h, 0, 0],
                    'rv_to_iv': res['irf'][h, 0, 1],
                    'iv_to_rv': res['irf'][h, 1, 0],
                    'rv_to_rv': res['irf'][h, 1, 1],
                })

    var_df = pd.DataFrame(var_results)
    var_df.to_csv(config.RESULTS_DIR / "var_results.csv", index=False)
    print(f"\n  VAR estimated for {len(var_results)} tickers")
    print(f"  Avg R² (IV equation): {var_df['r2_iv_eq'].mean():.4f}")
    print(f"  Avg R² (RV equation): {var_df['r2_rv_eq'].mean():.4f}")

    irf_df = pd.DataFrame(irf_all)
    irf_df.to_csv(config.RESULTS_DIR / "irf_results.csv", index=False)

    avg_irf = irf_df.groupby('horizon')[['iv_to_iv', 'rv_to_iv',
                                         'iv_to_rv', 'rv_to_rv']].mean()
    print("\n  Average Impulse Response Function:")
    print(f"  {'h':>3} {'IV→IV':>8} {'RV→IV':>8} {'IV→RV':>8} {'RV→RV':>8}")
    for h in range(0, IRF_PERIODS, 2):
        row = avg_irf.loc[h]
        print(f"  {h:>3} {row['iv_to_iv']:>8.4f} {row['rv_to_iv']:>8.4f} "
              f"{row['iv_to_rv']:>8.4f} {row['rv_to_rv']:>8.4f}")

    # ── C. Forecast-error variance decomposition ──
    print("\n--- C. FORECAST ERROR VARIANCE DECOMPOSITION ---")
    fevd_results = []
    for h in range(1, IRF_PERIODS + 1):
        sub_irf = irf_df[irf_df['horizon'] < h]
        avg = sub_irf.groupby('horizon')[['iv_to_rv', 'rv_to_rv']].mean()
        total_var = (avg['iv_to_rv'] ** 2 + avg['rv_to_rv'] ** 2).sum()
        iv_share = (avg['iv_to_rv'] ** 2).sum() / total_var if total_var > 0 else 0
        fevd_results.append({
            'horizon': h,
            'pct_rv_explained_by_iv': iv_share * 100,
            'pct_rv_explained_by_rv': (1 - iv_share) * 100,
        })

    fevd_df = pd.DataFrame(fevd_results)
    fevd_df.to_csv(config.RESULTS_DIR / "fevd_results.csv", index=False)
    print(f"  {'Horizon':>8} {'% RV by IV':>12} {'% RV by RV':>12}")
    for _, row in fevd_df.iterrows():
        print(f"  {int(row['horizon']):>8} {row['pct_rv_explained_by_iv']:>12.2f} "
              f"{row['pct_rv_explained_by_rv']:>12.2f}")

    print("\nGRANGER CAUSALITY & VAR ANALYSIS COMPLETE.")


if __name__ == "__main__":
    main()