#!/usr/bin/env python3
# =============================================================================
# Author: Simon-Pierre Boucher
# Contact: contact@spboucher.ai
# =============================================================================
"""Step 01 — Data extraction pipeline (raw → processed).

Extracts option-implied surface features and 5-minute realized volatility
from the raw DuckDB stores, then merges them into the master analysis panel.

Inputs (raw, external — see data/raw/README.md):
    options.duckdb, stock_5min.duckdb, etf_5min.duckdb, index_5min.duckdb

Outputs (data/processed/):
    options_features.parquet, realized_vol.parquet, merged_options_rv.parquet

This step is the only producer of the processed datasets; every other script
runs from its outputs. If the raw stores are unavailable the script exits
with a clear message and the shipped parquets remain authoritative.
"""

import sys
import warnings

import pandas as pd

import _bootstrap  # noqa: F401
from wp7 import config
from wp7.data_io import RawDataUnavailableError, open_raw_db, raw_db_path

warnings.filterwarnings('ignore')


# --------------------------------------------------------------------------
# 1A. Option-surface features per ticker per day
# --------------------------------------------------------------------------
def extract_options_features(con, tickers, batch_label=""):
    """Aggregate the daily option chain into ten implied-moment features."""
    ticker_str = ",".join([f"'{t}'" for t in tickers])
    query = f"""
    WITH base AS (
        SELECT ticker, trade_date, strike, expiry_date, call_put,
               bid_price, ask_price, bid_iv, ask_iv,
               open_interest, volume, delta, gamma, vega, theta, rho,
               (expiry_date - trade_date) AS dte,
               (bid_price + ask_price) / 2.0 AS mid_price,
               CASE WHEN bid_iv > 0 AND ask_iv > 0 THEN (bid_iv + ask_iv) / 2.0
                    WHEN ask_iv > 0 THEN ask_iv
                    ELSE bid_iv END AS mid_iv
        FROM option_chain
        WHERE ticker IN ({ticker_str})
          AND (expiry_date - trade_date) BETWEEN 7 AND 180
          AND ask_price > 0
    )
    SELECT
        ticker,
        trade_date,

        -- Implied Volatility: ATM (|delta| closest to 0.5)
        AVG(CASE WHEN call_put='c' AND ABS(delta - 0.5) < 0.1 AND dte BETWEEN 20 AND 40
                 THEN mid_iv END) AS iv_atm_30d,
        AVG(CASE WHEN call_put='c' AND ABS(delta - 0.5) < 0.1 AND dte BETWEEN 80 AND 100
                 THEN mid_iv END) AS iv_atm_90d,

        -- IV Term Structure slope (90d - 30d)
        AVG(CASE WHEN call_put='c' AND ABS(delta - 0.5) < 0.1 AND dte BETWEEN 80 AND 100
                 THEN mid_iv END) -
        AVG(CASE WHEN call_put='c' AND ABS(delta - 0.5) < 0.1 AND dte BETWEEN 20 AND 40
                 THEN mid_iv END) AS iv_term_slope,

        -- Volatility Skew (25d put IV - 25d call IV, ~30d expiry)
        AVG(CASE WHEN call_put='p' AND ABS(delta + 0.25) < 0.07 AND dte BETWEEN 20 AND 40
                 THEN mid_iv END) -
        AVG(CASE WHEN call_put='c' AND ABS(delta - 0.25) < 0.07 AND dte BETWEEN 20 AND 40
                 THEN mid_iv END) AS iv_skew_25d,

        -- Implied Skewness proxy (OTM put IV - OTM call IV normalized by ATM)
        (AVG(CASE WHEN call_put='p' AND ABS(delta + 0.10) < 0.05 AND dte BETWEEN 20 AND 40
                  THEN mid_iv END) -
         AVG(CASE WHEN call_put='c' AND ABS(delta - 0.10) < 0.05 AND dte BETWEEN 20 AND 40
                  THEN mid_iv END)) /
        NULLIF(AVG(CASE WHEN call_put='c' AND ABS(delta - 0.5) < 0.1 AND dte BETWEEN 20 AND 40
                        THEN mid_iv END), 0) AS implied_skewness,

        -- Implied Kurtosis proxy (wing IV avg / ATM IV)
        (AVG(CASE WHEN ABS(delta) < 0.15 AND ABS(delta) > 0.03 AND dte BETWEEN 20 AND 40
                  THEN mid_iv END)) /
        NULLIF(AVG(CASE WHEN call_put='c' AND ABS(delta - 0.5) < 0.1 AND dte BETWEEN 20 AND 40
                        THEN mid_iv END), 0) AS implied_kurtosis_proxy,

        -- Put-Call Volume Ratio
        SUM(CASE WHEN call_put='p' THEN volume ELSE 0 END)::DOUBLE /
        NULLIF(SUM(CASE WHEN call_put='c' THEN volume ELSE 0 END), 0) AS pc_volume_ratio,

        -- Put-Call OI Ratio
        SUM(CASE WHEN call_put='p' THEN open_interest ELSE 0 END)::DOUBLE /
        NULLIF(SUM(CASE WHEN call_put='c' THEN open_interest ELSE 0 END), 0) AS pc_oi_ratio,

        -- Aggregate Greeks
        SUM(CASE WHEN call_put='c' THEN gamma * open_interest ELSE 0 END) -
        SUM(CASE WHEN call_put='p' THEN gamma * open_interest ELSE 0 END) AS net_gamma_exposure,

        AVG(CASE WHEN dte BETWEEN 20 AND 40 THEN vega END) AS avg_vega_30d,
        AVG(CASE WHEN dte BETWEEN 20 AND 40 THEN theta END) AS avg_theta_30d,

        -- Total volume and OI
        SUM(volume) AS total_option_volume,
        SUM(open_interest) AS total_oi,

        -- Max OI strike (price magnet proxy)
        (SELECT b2.strike FROM base b2
         WHERE b2.ticker = base.ticker AND b2.trade_date = base.trade_date
         GROUP BY b2.strike ORDER BY SUM(b2.open_interest) DESC LIMIT 1) AS max_oi_strike

    FROM base
    GROUP BY ticker, trade_date
    HAVING COUNT(*) >= 10
    ORDER BY ticker, trade_date
    """
    print(f"  Extracting options features for {batch_label} ({len(tickers)} tickers)...")
    df = con.execute(query).fetchdf()
    print(f"  -> {len(df):,} rows extracted")
    return df


# --------------------------------------------------------------------------
# 1B. Realized volatility from 5-minute OHLCV
# --------------------------------------------------------------------------
def compute_realized_vol(db_name, table_name, id_col, tickers, label=""):
    """Daily realized variance, skewness and kurtosis from 5-minute bars."""
    ticker_str = ",".join([f"'{t}'" for t in tickers])
    con = open_raw_db(db_name)

    query = f"""
    WITH bars AS (
        SELECT {id_col} AS ticker,
               CAST(datetime AS DATE) AS trade_date,
               datetime,
               close,
               LAG(close) OVER (PARTITION BY {id_col} ORDER BY datetime) AS prev_close
        FROM {table_name}
        WHERE {id_col} IN ({ticker_str})
    ),
    returns AS (
        SELECT ticker, trade_date, datetime,
               LN(close / NULLIF(prev_close, 0)) AS log_ret
        FROM bars
        WHERE prev_close IS NOT NULL AND prev_close > 0 AND close > 0
    )
    SELECT
        ticker,
        trade_date,
        SUM(log_ret * log_ret) AS rv_daily,
        SQRT(SUM(log_ret * log_ret)) AS rvol_daily,
        COUNT(*) AS n_obs,
        FIRST(log_ret) AS open_ret,
        SUM(log_ret) AS daily_return,
        (SQRT(COUNT(*)) * SUM(POWER(log_ret, 3))) /
            NULLIF(POWER(SUM(log_ret * log_ret), 1.5), 0) AS realized_skew,
        (COUNT(*) * SUM(POWER(log_ret, 4))) /
            NULLIF(POWER(SUM(log_ret * log_ret), 2), 0) AS realized_kurt
    FROM returns
    GROUP BY ticker, trade_date
    HAVING COUNT(*) >= 20
    ORDER BY ticker, trade_date
    """
    print(f"  Computing realized vol for {label} ({len(tickers)} tickers)...")
    df = con.execute(query).fetchdf()
    print(f"  -> {len(df):,} rows")
    con.close()
    return df


# --------------------------------------------------------------------------
# 1C. Weekly aggregates and forward targets
# --------------------------------------------------------------------------
def compute_weekly_rv(rv_df):
    """Rolling weekly RV plus forward return / forward RV targets."""
    rv_df = rv_df.sort_values(['ticker', 'trade_date'])
    rv_df['rv_weekly'] = rv_df.groupby('ticker')['rv_daily'].transform(
        lambda x: x.rolling(5, min_periods=3).sum()
    )
    rv_df['ret_1d'] = rv_df.groupby('ticker')['daily_return'].shift(-1)
    rv_df['ret_5d'] = rv_df.groupby('ticker')['daily_return'].transform(
        lambda x: x.shift(-1).rolling(5, min_periods=3).sum()
    )
    rv_df['rv_fwd_1d'] = rv_df.groupby('ticker')['rv_daily'].shift(-1)
    rv_df['rv_fwd_5d'] = rv_df.groupby('ticker')['rv_daily'].transform(
        lambda x: x.shift(-1).rolling(5, min_periods=3).sum()
    )
    return rv_df


# --------------------------------------------------------------------------
# 1D. HAR-RV components
# --------------------------------------------------------------------------
def compute_har_components(rv_df):
    """Daily lag, weekly mean and monthly mean of realized variance."""
    rv_df = rv_df.sort_values(['ticker', 'trade_date'])
    rv_df['rv_lag1'] = rv_df.groupby('ticker')['rv_daily'].shift(1)
    rv_df['rv_w'] = rv_df.groupby('ticker')['rv_daily'].transform(
        lambda x: x.rolling(5, min_periods=3).mean()
    )
    rv_df['rv_m'] = rv_df.groupby('ticker')['rv_daily'].transform(
        lambda x: x.rolling(22, min_periods=10).mean()
    )
    return rv_df


def main():
    print("=" * 70)
    print("STEP 1: EXTRACTING OPTIONS-IMPLIED MOMENTS")
    print("=" * 70)
    config.ensure_output_dirs()

    # Options features
    opt_con = open_raw_db("options")
    opt_stocks = extract_options_features(opt_con, config.MAJOR_TICKERS, "Major Stocks")
    opt_etfs = extract_options_features(opt_con, config.ETF_TICKERS, "ETFs")
    opt_idx = extract_options_features(opt_con, config.INDEX_OPTION_TICKERS, "Indices")
    opt_con.close()

    opt_all = pd.concat([opt_stocks, opt_etfs, opt_idx], ignore_index=True)
    opt_all.to_parquet(config.OPTIONS_FEATURES_PARQUET, index=False)
    print(f"\nOptions features saved: {len(opt_all):,} rows")

    # Realized volatility per asset class
    rv_stocks = compute_realized_vol("stocks_5min", "ohlcv", "symbol",
                                     config.MAJOR_TICKERS, "Stocks 5min")
    rv_etfs = compute_realized_vol("etfs_5min", "ohlcv", "symbol",
                                   config.ETF_TICKERS, "ETFs 5min")
    rv_idx = compute_realized_vol("indices_5min", "ohlcv", "symbol",
                                  config.INDEX_SYMBOLS, "Indices 5min")

    rv_all = pd.concat([rv_stocks, rv_etfs, rv_idx], ignore_index=True)
    rv_all = compute_weekly_rv(rv_all)
    rv_all = compute_har_components(rv_all)
    rv_all.to_parquet(config.REALIZED_VOL_PARQUET, index=False)
    print(f"Realized vol saved: {len(rv_all):,} rows")

    # Merge options + realized vol into the master panel
    merged = pd.merge(opt_all, rv_all, on=['ticker', 'trade_date'], how='inner')
    merged = merged.sort_values(['ticker', 'trade_date']).reset_index(drop=True)
    merged.to_parquet(config.MERGED_PARQUET, index=False)
    print(f"\nMerged dataset saved: {len(merged):,} rows")
    print(f"Tickers: {merged['ticker'].nunique()}")
    print(f"Date range: {merged['trade_date'].min()} to {merged['trade_date'].max()}")
    print(f"Columns: {list(merged.columns)}")
    print("\nStep 1 COMPLETE.")


if __name__ == "__main__":
    try:
        main()
    except RawDataUnavailableError as exc:
        print(f"\n[SKIPPED] {exc}", file=sys.stderr)
        print(f"(expected stores: {[str(raw_db_path(n)) for n in config.RAW_DB_FILES]})",
              file=sys.stderr)
        sys.exit(2)