#!/usr/bin/env python3
# =============================================================================
# Author: Simon-Pierre Boucher
# Contact: contact@spboucher.ai
# =============================================================================
"""Step 02 — RQ1: Do option-implied moments predict next-day/next-week returns?

Pooled OLS (HC1) and Fama-MacBeth panel regressions across asset groups
(stocks, ETFs, indices, all).

Inputs : data/processed/merged_options_rv.parquet
Outputs: results/rq1_regression_results.csv, results/rq1_meta.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 fama_macbeth, winsorize

warnings.filterwarnings('ignore')

TARGETS = ['ret_1d', 'ret_5d']


def panel_ols(data, features, target, label=""):
    """Pooled OLS with HC1 t-statistics, as specified in the original study.

    Historical quirks preserved on purpose (they define the published
    numbers): features are standardized with the *pandas* sample std
    (ddof=1) and missing standardized cells are zero-filled; the reported
    ``p_value`` column is an ad-hoc normal-tail approximation, not an exact
    two-sided p-value (significance flags use the usual 1.96/2.576 cutoffs).
    """
    sub = data[features + [target, 'ticker', 'trade_date']].dropna()
    if len(sub) < 100:
        return None

    for f in features:
        sub[f] = winsorize(sub[f])
    sub[target] = winsorize(sub[target])

    means = sub[features].mean()
    stds = sub[features].std()
    X = (sub[features] - means) / stds
    X = X.fillna(0)
    X.insert(0, 'const', 1.0)
    y = sub[target].values

    coefs, _, _, _ = np.linalg.lstsq(X.values, y, rcond=None)

    y_pred = X.values @ coefs
    ss_res = np.sum((y - y_pred) ** 2)
    ss_tot = np.sum((y - y.mean()) ** 2)
    r2 = 1 - ss_res / ss_tot if ss_tot > 0 else 0

    n = len(y)
    k = len(features)
    adj_r2 = 1 - (1 - r2) * (n - 1) / (n - k - 1)

    # HC1 sandwich (X' diag(e²) X via broadcasting)
    e = y - y_pred
    XtX_inv = np.linalg.inv(X.values.T @ X.values)
    S = (X.values * (e ** 2)[:, None]).T @ X.values * n / (n - k - 1)
    se = np.sqrt(np.diag(XtX_inv @ S @ XtX_inv))
    t_stats = coefs / se

    results = pd.DataFrame({
        'variable': ['const'] + features,
        'coefficient': coefs,
        'std_error': se,
        't_stat': t_stats,
        'p_value': 2 * (1 - pd.Series(np.abs(t_stats)).apply(
            lambda x: min(1.0, 0.5 * np.exp(-0.5 * x**2) * np.sqrt(2 / np.pi) if x < 30 else 0)
        )).values,
        'significant_5pct': np.abs(t_stats) > 1.96,
        'significant_1pct': np.abs(t_stats) > 2.576,
    })

    meta = {
        'label': label,
        'target': target,
        'n_obs': n,
        'n_tickers': sub['ticker'].nunique(),
        'r_squared': r2,
        'adj_r_squared': adj_r2,
    }
    return results, meta


def main():
    print("=" * 70)
    print("RQ1: OPTION-IMPLIED MOMENTS AND RETURN PREDICTABILITY")
    print("=" * 70)
    config.ensure_output_dirs()

    df = load_merged()
    print(f"Loaded {len(df):,} rows, {df['ticker'].nunique()} tickers")

    stocks = df[~df['ticker'].isin(config.NON_STOCK_TICKERS)].copy()
    etfs = df[df['ticker'].isin(config.ETF_TICKERS)].copy()
    indices = df[df['ticker'].isin(config.INDEX_OPTION_TICKERS)].copy()

    features = config.RQ1_FEATURES
    all_results, all_meta = [], []

    for group_name, group_data in [("Stocks", stocks), ("ETFs", etfs),
                                   ("Indices", indices), ("All", df)]:
        for target in TARGETS:
            horizon = "1-Day" if target == 'ret_1d' else "5-Day"
            label = f"{group_name} | {horizon}"

            # Pooled OLS
            res = panel_ols(group_data, features, target, label)
            if res:
                r, m = res
                r['group'] = group_name
                r['horizon'] = horizon
                r['method'] = 'Pooled OLS'
                all_results.append(r)
                all_meta.append(m)
                print(f"\n{label} (Pooled OLS): R²={m['r_squared']:.6f}, "
                      f"Adj-R²={m['adj_r_squared']:.6f}, N={m['n_obs']:,}")
                sig = r[r['significant_5pct'] & (r['variable'] != 'const')]
                if len(sig) > 0:
                    print(f"  Significant predictors: {', '.join(sig['variable'].tolist())}")

            # Fama-MacBeth (stocks and full panel only)
            if group_name in ['Stocks', 'All']:
                res_fm = fama_macbeth(group_data, features, target)
                if res_fm:
                    r_fm, n_periods, n_tickers = res_fm
                    r_fm['group'] = group_name
                    r_fm['horizon'] = horizon
                    r_fm['method'] = 'Fama-MacBeth'
                    all_results.append(r_fm)
                    all_meta.append({'label': label, 'target': target,
                                     'n_periods': n_periods, 'n_tickers': n_tickers})
                    print(f"  Fama-MacBeth: N_periods={n_periods}")
                    sig_fm = r_fm[r_fm['fm_significant_5pct'] & (r_fm['variable'] != 'const')]
                    if len(sig_fm) > 0:
                        print(f"  FM Significant: {', '.join(sig_fm['variable'].tolist())}")

    results_df = pd.concat(all_results, ignore_index=True)
    results_df.to_csv(config.RESULTS_DIR / "rq1_regression_results.csv", index=False)

    meta_df = pd.DataFrame(all_meta)
    meta_df.to_csv(config.RESULTS_DIR / "rq1_meta.csv", index=False)

    print("\n" + "=" * 70)
    print("RQ1 SUMMARY TABLE")
    print("=" * 70)
    summary_rows = [
        {'Group': m.get('label', ''),
         'N': m.get('n_obs', m.get('n_periods', '')),
         'R²': f"{m.get('r_squared', 0):.6f}",
         'Adj-R²': f"{m.get('adj_r_squared', 0):.6f}"}
        for _, m in meta_df.iterrows() if 'r_squared' in m
    ]
    if summary_rows:
        print(pd.DataFrame(summary_rows).to_string(index=False))

    print("\nRQ1 COMPLETE.")


if __name__ == "__main__":
    main()