"""
Volatility Analysis Module
==========================
Computes realized volatility, implied volatility analysis,
term structure, skew metrics, and volatility regime detection.

Functions:
  - realized_vol: Historical/realized volatility
  - parkinson_vol: Parkinson high-low volatility estimator
  - garman_klass_vol: Garman-Klass OHLC volatility estimator
  - vol_cone: Volatility cone (percentiles over rolling windows)
  - iv_term_structure: Implied vol term structure from chain
  - iv_skew_metrics: Skew and smile metrics
  - vol_regime: Volatility regime detection
  - vol_risk_premium: IV - RV spread (vol risk premium)
  - vix_fair_value: VIX fair value from options chain
  - term_structure_slope: Contango/backwardation slope
"""

import numpy as np
import pandas as pd
from typing import Optional, Dict, List, Any, Tuple
from scipy import stats


# ── Realized Volatility ─────────────────────────────────────────────────────

def realized_vol(
    prices: pd.Series,
    window: int = 20,
    annualize: bool = True
) -> pd.Series:
    """
    Compute rolling realized volatility from price series.
    
    Parameters
    ----------
    prices : Price series (close prices)
    window : Rolling window (trading days)
    annualize : If True, multiply by sqrt(252)
    
    Returns
    -------
    Series of realized volatility (decimal)
    """
    if prices is None or len(prices) < 2:
        return pd.Series(dtype=float)
    
    log_returns = np.log(prices / prices.shift(1))
    vol = log_returns.rolling(window=window).std()
    
    if annualize:
        vol = vol * np.sqrt(252)
    
    return vol


def parkinson_vol(
    high: pd.Series,
    low: pd.Series,
    window: int = 20,
    annualize: bool = True
) -> pd.Series:
    """
    Parkinson (1980) volatility estimator using high-low range.
    
    More efficient than close-to-close estimator.
    sigma^2 = (1 / (4 * N * ln(2))) * sum(ln(Hi/Li)^2)
    """
    if high is None or low is None or len(high) < 2:
        return pd.Series(dtype=float)
    
    log_hl = np.log(high / low)
    factor = 1.0 / (4.0 * np.log(2))
    
    vol_sq = (log_hl ** 2).rolling(window=window).mean() * factor
    vol = np.sqrt(vol_sq)
    
    if annualize:
        vol = vol * np.sqrt(252)
    
    return vol


def garman_klass_vol(
    open_: pd.Series,
    high: pd.Series,
    low: pd.Series,
    close: pd.Series,
    window: int = 20,
    annualize: bool = True
) -> pd.Series:
    """
    Garman-Klass (1980) volatility estimator using OHLC.
    
    More efficient than Parkinson (uses open and close).
    sigma^2 = 0.5 * ln(H/L)^2 - (2*ln(2)-1) * ln(C/O)^2
    """
    if any(s is None for s in [open_, high, low, close]):
        return pd.Series(dtype=float)
    
    log_hl = np.log(high / low)
    log_co = np.log(close / open_)
    
    vol_sq = (0.5 * log_hl**2 - (2 * np.log(2) - 1) * log_co**2).rolling(window=window).mean()
    vol = np.sqrt(vol_sq.clip(lower=0))  # Clip negative values
    
    if annualize:
        vol = vol * np.sqrt(252)
    
    return vol


def rogers_satchell_vol(
    open_: pd.Series,
    high: pd.Series,
    low: pd.Series,
    close: pd.Series,
    window: int = 20,
    annualize: bool = True
) -> pd.Series:
    """
    Rogers-Satchell (1991) volatility estimator.
    
    Works with drifting assets (unlike Garman-Klass).
    """
    if any(s is None for s in [open_, high, low, close]):
        return pd.Series(dtype=float)
    
    log_ho = np.log(high / open_)
    log_hc = np.log(high / close)
    log_lo = np.log(low / open_)
    log_lc = np.log(low / close)
    
    vol_sq = (log_ho * log_hc + log_lo * log_lc).rolling(window=window).mean()
    vol = np.sqrt(vol_sq.clip(lower=0))
    
    if annualize:
        vol = vol * np.sqrt(252)
    
    return vol


# ── Volatility Cone ─────────────────────────────────────────────────────────

def vol_cone(
    prices: pd.Series,
    windows: List[int] = None,
    percentiles: List[float] = None,
) -> pd.DataFrame:
    """
    Compute volatility cone.
    
    Shows realized volatility percentiles across different time horizons.
    Useful for determining if current IV is rich/cheap vs historical RV.
    
    Parameters
    ----------
    prices : Price series
    windows : List of lookback windows (default: [5, 10, 20, 30, 60, 90, 120, 252])
    percentiles : Percentiles to compute (default: [5, 25, 50, 75, 95])
    
    Returns
    -------
    DataFrame with windows as rows, percentiles as columns
    """
    if prices is None or len(prices) < 252:
        return pd.DataFrame()
    
    if windows is None:
        windows = [5, 10, 20, 30, 60, 90, 120, 252]
    if percentiles is None:
        percentiles = [5, 25, 50, 75, 95]
    
    log_returns = np.log(prices / prices.shift(1)).dropna()
    
    data = {}
    for w in windows:
        if len(log_returns) < w:
            continue
        # Rolling vol for each possible window
        rolling_vols = []
        for i in range(w, len(log_returns)):
            vol = log_returns.iloc[i-w:i].std() * np.sqrt(252)
            rolling_vols.append(vol)
        
        if rolling_vols:
            row = {}
            for p in percentiles:
                row[f'p{p}'] = np.percentile(rolling_vols, p)
            data[w] = row
    
    if not data:
        return pd.DataFrame()
    
    df = pd.DataFrame(data).T
    df.index.name = 'window'
    return df


# ── IV Term Structure ──────────────────────────────────────────────────────

def iv_term_structure(
    chain_records: List[Dict[str, Any]],
    spot: float
) -> pd.DataFrame:
    """
    Compute implied volatility term structure from options chain.
    
    Groups by expiry and computes average ATM IV for each tenor.
    
    Returns DataFrame with columns:
      expiry, dte, atm_iv_call, atm_iv_put, avg_iv, skew
    """
    if not chain_records:
        return pd.DataFrame()
    
    # Group by expiry
    by_expiry: Dict[str, List[Dict]] = {}
    for rec in chain_records:
        exp = rec.get('expiry', 'unknown')
        if exp not in by_expiry:
            by_expiry[exp] = []
        by_expiry[exp].append(rec)
    
    rows = []
    for exp, records in sorted(by_expiry.items()):
        # Find ATM records (closest to spot)
        atm_records = sorted(records, key=lambda r: abs(r.get('strike', 0) - spot))[:4]
        
        if not atm_records:
            continue
        
        iv_calls = [r.get('iv_call', 0) for r in atm_records if r.get('iv_call', 0) > 0]
        iv_puts = [r.get('iv_put', 0) for r in atm_records if r.get('iv_put', 0) > 0]
        
        atm_iv_c = np.mean(iv_calls) if iv_calls else 0
        atm_iv_p = np.mean(iv_puts) if iv_puts else 0
        avg_iv = (atm_iv_c + atm_iv_p) / 2.0
        skew = atm_iv_p - atm_iv_c
        
        # Compute DTE
        try:
            exp_date = pd.to_datetime(exp)
            dte = max(0, (exp_date - pd.Timestamp.now()).days)
        except:
            dte = 0
        
        rows.append({
            'expiry': exp,
            'dte': dte,
            'atm_iv_call': atm_iv_c,
            'atm_iv_put': atm_iv_p,
            'avg_iv': avg_iv,
            'skew': skew,
        })
    
    return pd.DataFrame(rows).sort_values('dte').reset_index(drop=True)


def term_structure_slope(term_struct_df: pd.DataFrame) -> Dict[str, float]:
    """
    Compute term structure slope metrics.
    
    Returns dict with:
      - front_month_iv: IV of nearest expiry
      - back_month_iv: IV of furthest expiry
      - slope: Back - Front (positive = contango)
      - slope_pct: Slope as percentage of front month
      - structure: 'Contango' or 'Backwardation'
    """
    if term_struct_df.empty or len(term_struct_df) < 2:
        return {
            'front_month_iv': 0, 'back_month_iv': 0,
            'slope': 0, 'slope_pct': 0, 'structure': 'N/A',
        }
    
    front = term_struct_df.iloc[0]['avg_iv']
    back = term_struct_df.iloc[-1]['avg_iv']
    slope = back - front
    slope_pct = (slope / front * 100) if front > 0 else 0
    
    return {
        'front_month_iv': front,
        'back_month_iv': back,
        'slope': slope,
        'slope_pct': slope_pct,
        'structure': 'Contango' if slope > 0 else 'Backwardation',
    }


# ── Skew Metrics ────────────────────────────────────────────────────────────

def compute_iv_skew(
    chain_records: List[Dict[str, Any]],
    spot: float
) -> Dict[str, Any]:
    """
    Compute comprehensive IV skew metrics.
    
    Returns dict with:
      - raw_skew: Put IV - Call IV at same strike (average)
      - delta_skew: 25d risk reversal
      - butterfly: 25d butterfly
      - smile_curvature: Second derivative of IV smile
      - skew_slope: Slope of skew (IV change per 1% moneyness)
    """
    if not chain_records:
        return {
            'raw_skew': 0, 'delta_skew': 0, 'butterfly': 0,
            'smile_curvature': 0, 'skew_slope': 0,
        }
    
    # Compute moneyness and skew for each record
    data = []
    for rec in chain_records:
        strike = rec.get('strike', 0)
        if strike <= 0:
            continue
        iv_c = rec.get('iv_call', 0)
        iv_p = rec.get('iv_put', 0)
        moneyness = (strike / spot - 1) * 100  # Percent
        
        data.append({
            'strike': strike,
            'moneyness': moneyness,
            'iv_call': iv_c,
            'iv_put': iv_p,
            'raw_skew': iv_p - iv_c,
        })
    
    if not data:
        return {
            'raw_skew': 0, 'delta_skew': 0, 'butterfly': 0,
            'smile_curvature': 0, 'skew_slope': 0,
        }
    
    df = pd.DataFrame(data)
    
    # Raw skew: average put-call IV difference
    raw_skew = df['raw_skew'].mean()
    
    # Skew slope: regression of raw skew on moneyness
    if len(df) > 2:
        slope, intercept, r_value, p_value, std_err = stats.linregress(
            df['moneyness'], df['raw_skew']
        )
        skew_slope = slope
    else:
        skew_slope = 0
    
    # ATM skew (at moneyness ~ 0)
    atm_records = df[df['moneyness'].abs() < 1.0]
    atm_skew = atm_records['raw_skew'].mean() if len(atm_records) > 0 else raw_skew
    
    # Wing skew (OTM puts vs OTM calls)
    otm_puts = df[df['moneyness'] < -2]
    otm_calls = df[df['moneyness'] > 2]
    wing_skew = (
        otm_puts['iv_put'].mean() - otm_calls['iv_call'].mean()
        if len(otm_puts) > 0 and len(otm_calls) > 0
        else 0
    )
    
    return {
        'raw_skew': raw_skew,
        'atm_skew': atm_skew,
        'wing_skew': wing_skew,
        'skew_slope': skew_slope,
        'delta_skew': wing_skew,  # Approximation
        'butterfly': atm_skew,  # Simplified
    }


# ── Volatility Regime ───────────────────────────────────────────────────────

def vol_regime(
    current_iv: float,
    rv_20d: float,
    rv_60d: float,
    iv_rank: float,
    iv_percentile: float,
    term_slope: float
) -> Dict[str, Any]:
    """
    Determine current volatility regime.
    
    Parameters
    ----------
    current_iv : Current implied vol (e.g., VIX)
    rv_20d : 20-day realized vol
    rv_60d : 60-day realized vol
    iv_rank : IV rank (0-100)
    iv_percentile : IV percentile (0-100)
    term_slope : Term structure slope (back - front)
    
    Returns
    -------
    Dict with regime classification and signals.
    """
    # Regime classification
    if iv_rank > 75:
        regime = 'High Vol'
        regime_color = 'red'
    elif iv_rank > 50:
        regime = 'Above Normal'
        regime_color = 'orange'
    elif iv_rank > 25:
        regime = 'Below Normal'
        regime_color = 'yellow'
    else:
        regime = 'Low Vol'
        regime_color = 'green'
    
    # Vol risk premium
    vrp = current_iv - rv_20d
    
    # Term structure signal
    if term_slope > 2:
        term_signal = 'Strong Contango'
    elif term_slope > 0.5:
        term_signal = 'Contango'
    elif term_slope > -0.5:
        term_signal = 'Flat'
    elif term_slope > -2:
        term_signal = 'Backwardation'
    else:
        term_signal = 'Strong Backwardation'
    
    # RV trend
    if rv_20d > rv_60d * 1.2:
        rv_trend = 'Rising'
    elif rv_20d < rv_60d * 0.8:
        rv_trend = 'Falling'
    else:
        rv_trend = 'Stable'
    
    # Composite signal
    signals = []
    if iv_rank > 80:
        signals.append('IV Elevated - Consider selling vol')
    elif iv_rank < 20:
        signals.append('IV Compressed - Consider buying vol')
    
    if vrp > 5:
        signals.append('High Vol Risk Premium - Favor vol selling')
    elif vrp < -2:
        signals.append('Negative VRP - Favor vol buying')
    
    if 'Backwardation' in term_signal:
        signals.append('Term Structure Inverted - Stress signal')
    
    if rv_trend == 'Rising' and iv_rank < 50:
        signals.append('RV rising but IV not catching up - Potential vol buy')
    
    return {
        'regime': regime,
        'regime_color': regime_color,
        'iv_rank': iv_rank,
        'iv_percentile': iv_percentile,
        'vrp': vrp,
        'term_signal': term_signal,
        'rv_trend': rv_trend,
        'signals': signals,
    }


def compute_iv_rank(iv_series: pd.Series, lookback: int = 252) -> float:
    """
    Compute IV Rank: (Current IV - 52w Low) / (52w High - 52w Low) * 100
    
    IV Rank measures where current IV sits in its 52-week range.
    """
    if iv_series is None or len(iv_series) < 20:
        return 50.0  # Neutral default
    
    recent = iv_series.tail(lookback)
    current = iv_series.iloc[-1]
    low = recent.min()
    high = recent.max()
    
    if high <= low:
        return 50.0
    
    return float((current - low) / (high - low) * 100)


def compute_iv_percentile(iv_series: pd.Series, lookback: int = 252) -> float:
    """
    Compute IV Percentile: % of days in lookback where IV was below current.
    
    More robust than IV Rank (not affected by single outlier).
    """
    if iv_series is None or len(iv_series) < 20:
        return 50.0
    
    recent = iv_series.tail(lookback)
    current = iv_series.iloc[-1]
    below = (recent < current).sum()
    
    return float(below / len(recent) * 100)


# ── VIX Fair Value ──────────────────────────────────────────────────────────

def vix_fair_value(
    chain_records: List[Dict[str, Any]],
    spot: float,
    dte_target: int = 30
) -> Dict[str, float]:
    """
    Estimate VIX fair value from SPX options chain.
    
    Uses the VIX methodology: weighted sum of OTM option prices
    to estimate 30-day implied variance.
    
    Returns dict with:
      - vix_fair: Estimated VIX fair value
      - vix_30d: 30-day expiry IV (if available)
      - variance_30d: 30-day variance
    """
    if not chain_records:
        return {'vix_fair': 0, 'vix_30d': 0, 'variance_30d': 0}
    
    # Find closest expiry to 30 DTE
    by_expiry: Dict[str, List[Dict]] = {}
    for rec in chain_records:
        exp = rec.get('expiry', '')
        if exp not in by_expiry:
            by_expiry[exp] = []
        by_expiry[exp].append(rec)
    
    best_expiry = None
    best_dte_diff = 999
    
    for exp in by_expiry:
        try:
            exp_date = pd.to_datetime(exp)
            dte = (exp_date - pd.Timestamp.now()).days
            if abs(dte - dte_target) < best_dte_diff:
                best_dte_diff = abs(dte - dte_target)
                best_expiry = exp
        except:
            continue
    
    if best_expiry is None:
        return {'vix_fair': 0, 'vix_30d': 0, 'variance_30d': 0}
    
    records = by_expiry[best_expiry]
    
    # Compute ATM IV as proxy for VIX
    atm_records = sorted(records, key=lambda r: abs(r.get('strike', 0) - spot))[:6]
    ivs = []
    for r in atm_records:
        iv_c = r.get('iv_call', 0)
        iv_p = r.get('iv_put', 0)
        if iv_c > 0:
            ivs.append(iv_c)
        if iv_p > 0:
            ivs.append(iv_p)
    
    atm_iv = np.mean(ivs) if ivs else 0.20
    
    # Simple VIX approximation: ATM IV * 100
    vix_fair = atm_iv * 100
    
    return {
        'vix_fair': vix_fair,
        'vix_30d': atm_iv * 100,
        'variance_30d': atm_iv ** 2,
        'expiry_used': best_expiry,
        'dte_used': dte_target,
    }


# ── Vol Risk Premium ────────────────────────────────────────────────────────

def vol_risk_premium(
    iv_series: pd.Series,
    rv_series: pd.Series,
    window: int = 20
) -> Dict[str, float]:
    """
    Compute volatility risk premium (IV - RV).
    
    Positive VRP = Implied vol > Realized vol (normal, vol sellers win on avg)
    Negative VRP = Realized vol > Implied vol (rare, stress periods)
    
    Returns dict with:
      - current_vrp: Current IV - RV
      - avg_vrp: Average VRP over window
      - vrp_zscore: Z-score of current VRP
      - vrp_regime: 'High', 'Normal', 'Low', 'Negative'
    """
    if iv_series is None or rv_series is None:
        return {'current_vrp': 0, 'avg_vrp': 0, 'vrp_zscore': 0, 'vrp_regime': 'N/A'}
    
    # Align series
    common_idx = iv_series.index.intersection(rv_series.index)
    if len(common_idx) < 10:
        return {'current_vrp': 0, 'avg_vrp': 0, 'vrp_zscore': 0, 'vrp_regime': 'N/A'}
    
    iv = iv_series.loc[common_idx]
    rv = rv_series.loc[common_idx]
    
    vrp = iv - rv
    current_vrp = vrp.iloc[-1]
    avg_vrp = vrp.tail(window).mean()
    std_vrp = vrp.tail(window).std()
    
    zscore = (current_vrp - avg_vrp) / std_vrp if std_vrp > 0 else 0
    
    if current_vrp < 0:
        regime = 'Negative'
    elif zscore > 1.5:
        regime = 'High'
    elif zscore < -1.0:
        regime = 'Low'
    else:
        regime = 'Normal'
    
    return {
        'current_vrp': float(current_vrp),
        'avg_vrp': float(avg_vrp),
        'vrp_zscore': float(zscore),
        'vrp_regime': regime,
    }


# ── Summary ─────────────────────────────────────────────────────────────────

def vol_summary(
    prices: pd.Series,
    current_iv: float,
    chain_records: Optional[List[Dict[str, Any]]] = None,
    spot: Optional[float] = None,
) -> Dict[str, Any]:
    """
    Compute comprehensive volatility summary.
    
    Returns dict with all key vol metrics.
    """
    if prices is None or len(prices) < 60:
        return {'error': 'Insufficient price data'}
    
    # Realized vols
    rv_5d = realized_vol(prices, window=5).iloc[-1] * 100 if len(prices) > 5 else 0
    rv_20d = realized_vol(prices, window=20).iloc[-1] * 100 if len(prices) > 20 else 0
    rv_60d = realized_vol(prices, window=60).iloc[-1] * 100 if len(prices) > 60 else 0
    
    # IV rank/percentile (use current_iv as proxy if no series)
    iv_rank = 50.0  # Default
    iv_pct = 50.0
    
    # VRP
    vrp = current_iv - rv_20d if rv_20d > 0 else 0
    
    # Term structure
    term_slope = 0.0
    if chain_records and spot:
        ts = iv_term_structure(chain_records, spot)
        if not ts.empty and len(ts) >= 2:
            term_slope = ts.iloc[-1]['avg_iv'] - ts.iloc[0]['avg_iv']
    
    # Regime
    regime = vol_regime(current_iv, rv_20d, rv_60d, iv_rank, iv_pct, term_slope)
    
    return {
        'current_iv': current_iv,
        'rv_5d': rv_5d,
        'rv_20d': rv_20d,
        'rv_60d': rv_60d,
        'iv_rank': iv_rank,
        'iv_percentile': iv_pct,
        'vrp': vrp,
        'term_slope': term_slope,
        'regime': regime['regime'],
        'regime_color': regime['regime_color'],
        'signals': regime['signals'],
    }


if __name__ == '__main__':
    # Smoke test
    np.random.seed(42)
    n = 252
    returns = np.random.normal(0.0003, 0.012, n)
    prices = pd.Series(100 * np.exp(np.cumsum(returns)),
                       index=pd.date_range('2025-01-01', periods=n, freq='B'))
    
    print("Volatility Summary:")
    rv_20 = realized_vol(prices, window=20)
    rv_60 = realized_vol(prices, window=60)
    print(f"  RV 20d: {rv_20.iloc[-1]*100:.2f}%")
    print(f"  RV 60d: {rv_60.iloc[-1]*100:.2f}%")
    
    cone = vol_cone(prices)
    if not cone.empty:
        print(f"\nVol Cone (20d):")
        print(f"  5th %ile: {cone.loc[20, 'p5']*100:.2f}%")
        print(f"  Median:   {cone.loc[20, 'p50']*100:.2f}%")
        print(f"  95th %ile: {cone.loc[20, 'p95']*100:.2f}%")
    
    summary = vol_summary(prices, current_iv=18.5)
    print(f"\nFull Summary:")
    for k, v in summary.items():
        if isinstance(v, float):
            print(f"  {k}: {v:.2f}")
        elif isinstance(v, list):
            print(f"  {k}: {v}")
        else:
            print(f"  {k}: {v}")