#!/usr/bin/env python3
# =============================================================================
# Author: Simon-Pierre Boucher
# Contact: contact@spboucher.ai
# =============================================================================
"""Step 07 — Extended descriptive statistics and data-quality analysis.

Panels A–E and G run from the processed panel alone. Panel F (VIX regimes)
and Panel H (raw options data quality) additionally need the raw stores and
are skipped gracefully when those are absent.

Inputs : data/processed/merged_options_rv.parquet
         (+ index_5min.duckdb and options.duckdb for panels F/H)
Outputs: results/descriptive_*.csv (8 tables)
"""

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, open_raw_db

warnings.filterwarnings('ignore')

SUMMARY_VARS = [
    'iv_atm_30d', 'iv_atm_90d', 'iv_term_slope', 'iv_skew_25d',
    'implied_skewness', 'implied_kurtosis_proxy',
    'pc_volume_ratio', 'pc_oi_ratio', 'net_gamma_exposure',
    'avg_vega_30d', 'avg_theta_30d', 'total_option_volume', 'total_oi',
    'rv_daily', 'rvol_daily', 'rv_weekly', 'daily_return',
    'realized_skew', 'realized_kurt',
    'ret_1d', 'ret_5d', 'rv_fwd_1d', 'rv_fwd_5d',
]

CORR_VARS = ['iv_atm_30d', 'iv_term_slope', 'iv_skew_25d', 'implied_skewness',
             'implied_kurtosis_proxy', 'pc_volume_ratio', 'pc_oi_ratio',
             'rv_daily', 'rv_weekly', 'ret_1d', 'ret_5d']

AUTOCORR_VARS = ['iv_atm_30d', 'iv_skew_25d', 'rv_daily', 'daily_return', 'pc_volume_ratio']


def group_stats(data: pd.DataFrame, label: str) -> dict:
    """One summary row per asset group."""
    return {
        'Group': label,
        'N_obs': len(data),
        'N_tickers': data['ticker'].nunique(),
        'Date_min': str(data['trade_date'].min().date()),
        'Date_max': str(data['trade_date'].max().date()),
        'Mean_IV_ATM': data['iv_atm_30d'].mean(),
        'Std_IV_ATM': data['iv_atm_30d'].std(),
        'Mean_RV': data['rv_daily'].mean(),
        'Mean_Skew': data['iv_skew_25d'].mean(),
        'Mean_Ret_1d': data['ret_1d'].mean(),
        'Std_Ret_1d': data['ret_1d'].std(),
        'Mean_PC_ratio': data['pc_volume_ratio'].mean(),
    }


def main():
    print("=" * 70)
    print("EXTENDED DESCRIPTIVE STATISTICS")
    print("=" * 70)
    config.ensure_output_dirs()
    res_dir = config.RESULTS_DIR

    df = load_merged()
    df['year'] = df['trade_date'].dt.year

    # ── Panel A: summary statistics ──
    print("\n--- PANEL A: SUMMARY STATISTICS ---")
    summary = df[SUMMARY_VARS].describe(
        percentiles=[0.01, 0.05, 0.25, 0.5, 0.75, 0.95, 0.99]).T
    summary['skewness'] = df[SUMMARY_VARS].skew()
    summary['kurtosis'] = df[SUMMARY_VARS].kurtosis()
    summary['pct_missing'] = df[SUMMARY_VARS].isnull().mean() * 100
    summary.to_csv(res_dir / "descriptive_summary_stats.csv")
    print(summary[['count', 'mean', 'std', '1%', '50%', '99%',
                   'skewness', 'kurtosis', 'pct_missing']].round(4).to_string())

    # ── Panel B: coverage by year ──
    print("\n--- PANEL B: COVERAGE BY YEAR ---")
    coverage = df.groupby('year').agg(
        n_obs=('ticker', 'count'),
        n_tickers=('ticker', 'nunique'),
        avg_iv_atm=('iv_atm_30d', 'mean'),
        avg_rv=('rv_daily', 'mean'),
        avg_skew=('iv_skew_25d', 'mean'),
        avg_ret=('daily_return', 'mean'),
        std_ret=('daily_return', 'std'),
    ).round(6)
    coverage.to_csv(res_dir / "descriptive_coverage_by_year.csv")
    print(coverage.to_string())

    # ── Panel C: coverage by asset group ──
    print("\n--- PANEL C: BY ASSET GROUP ---")
    stocks_list = [t for t in df['ticker'].unique() if t not in config.NON_STOCK_TICKERS]
    groups = pd.DataFrame([
        group_stats(df[df['ticker'].isin(stocks_list)], 'Stocks'),
        group_stats(df[df['ticker'].isin(config.ETF_TICKERS)], 'ETFs'),
        group_stats(df[df['ticker'].isin(config.INDEX_OPTION_TICKERS)], 'Indices'),
        group_stats(df, 'All'),
    ])
    groups.to_csv(res_dir / "descriptive_by_group.csv", index=False)
    print(groups.to_string(index=False))

    # ── Panel D: correlation matrix ──
    print("\n--- PANEL D: CORRELATION MATRIX ---")
    corr_matrix = df[CORR_VARS].corr().round(3)
    corr_matrix.to_csv(res_dir / "descriptive_correlation_matrix.csv")
    print(corr_matrix.to_string())

    # ── Panel E: autocorrelation structure ──
    print("\n--- PANEL E: AUTOCORRELATION STRUCTURE ---")
    autocorr_results = []
    for var in AUTOCORR_VARS:
        for lag in [1, 5, 10, 22]:
            ac = df.groupby('ticker')[var].apply(lambda x: x.autocorr(lag=lag)).mean()
            autocorr_results.append({'variable': var, 'lag': lag, 'avg_autocorr': ac})
    autocorr_df = pd.DataFrame(autocorr_results)
    autocorr_df.to_csv(res_dir / "descriptive_autocorrelations.csv", index=False)
    print(autocorr_df.pivot(index='variable', columns='lag', values='avg_autocorr')
          .round(4).to_string())

    # ── Panel G: cross-sectional dispersion by year ──
    print("\n--- PANEL G: CROSS-SECTIONAL DISPERSION ---")
    cs_disp = df.groupby('year').agg(
        iv_atm_cs_std=('iv_atm_30d', 'std'),
        skew_cs_std=('iv_skew_25d', 'std'),
        rv_cs_std=('rv_daily', 'std'),
        ret_cs_std=('daily_return', 'std'),
        n_tickers=('ticker', 'nunique'),
    ).round(6)
    cs_disp.to_csv(res_dir / "descriptive_cross_sectional_dispersion.csv")
    print(cs_disp.to_string())

    # ── Panel F: statistics by VIX regime (needs raw VIX series) ──
    print("\n--- PANEL F: STATISTICS BY VIX REGIME ---")
    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_VERBOSE)
        regime_stats = dfv.groupby('vix_regime', observed=True).agg(
            n_obs=('ticker', 'count'),
            pct_obs=('ticker', lambda x: len(x) / len(dfv) * 100),
            mean_iv_atm=('iv_atm_30d', 'mean'),
            mean_rv=('rv_daily', 'mean'),
            mean_skew=('iv_skew_25d', 'mean'),
            mean_ret_1d=('ret_1d', 'mean'),
            std_ret_1d=('ret_1d', 'std'),
            mean_pc_ratio=('pc_volume_ratio', 'mean'),
            mean_impl_skew=('implied_skewness', 'mean'),
        ).round(6)
        regime_stats.to_csv(res_dir / "descriptive_vix_regimes.csv")
        print(regime_stats.to_string())
    except RawDataUnavailableError as exc:
        print(f"  [Panel F skipped — raw stores unavailable]\n  {exc}")

    # ── Panel H: raw options data quality (needs options.duckdb) ──
    print("\n--- PANEL H: OPTIONS DATA QUALITY ---")
    try:
        opt_con = open_raw_db("options")
        quality = opt_con.execute("""
            SELECT
                EXTRACT(YEAR FROM trade_date) AS year,
                COUNT(*) AS n_records,
                COUNT(DISTINCT ticker) AS n_tickers,
                AVG(CASE WHEN bid_iv > 0 AND ask_iv > 0 THEN 1.0 ELSE 0.0 END) AS pct_valid_iv,
                AVG(CASE WHEN volume > 0 THEN 1.0 ELSE 0.0 END) AS pct_with_volume,
                AVG(CASE WHEN open_interest > 0 THEN 1.0 ELSE 0.0 END) AS pct_with_oi,
                AVG(CASE WHEN delta IS NOT NULL AND delta != 0 THEN 1.0 ELSE 0.0 END) AS pct_valid_greeks,
                AVG(ask_price - bid_price) AS avg_spread
            FROM option_chain
            GROUP BY year
            ORDER BY year
        """).fetchdf()
        opt_con.close()
        quality.to_csv(res_dir / "descriptive_options_quality.csv", index=False)
        print(quality.round(4).to_string(index=False))
    except RawDataUnavailableError as exc:
        print(f"  [Panel H skipped — raw stores unavailable]\n  {exc}")

    print("\nDESCRIPTIVE STATISTICS COMPLETE.")


if __name__ == "__main__":
    main()