"""
Intraday Volatility Transmission Across Asset Classes
Diebold-Yilmaz Connectedness Analysis -- stage 1 (May 2026 pipeline)

NOTE (revision of 2026-10-17, see docs/CHANGELOG-REVISION.md):
  * This script is kept as the sole producer of results/daily_rv_all.csv and
    results/table1_descriptive_stats.csv from the vendor's 5-minute DuckDB bars.
  * Its sections 3-8 (VAR/GFEVD, rolling, regimes, sessions, Granger, network)
    are SUPERSEDED by scripts/02_connectedness_revision.py and
    scripts/03_intraday_sessions.py. Known defects of the original stages:
      - the VAR panel forward-filled weekends/holidays (rv_panel.ffill(limit=5)),
        so 1,137 of 3,651 rows contained stale values;
      - realized volatility entered the VAR in levels, not logs;
      - FirstRate time stamps are U.S. Eastern, not UTC (comment below is wrong),
        so the session labels of the original Fig. 5 were shifted by 4-5 hours;
      - the "session" Granger tests repeated the same daily test three times;
      - 5-minute returns beyond 20 sample standard deviations are dropped before
        computing RV (undisclosed in v1.0; disclosed in v1.1).
  The code below is left unchanged for the audit trail.
"""

import os
import warnings
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import matplotlib.gridspec as gridspec
import seaborn as sns
import duckdb
from scipy import stats
import networkx as nx
from statsmodels.tsa.api import VAR
from statsmodels.tsa.stattools import grangercausalitytests
import traceback

warnings.filterwarnings('ignore')

# ─── Paths ────────────────────────────────────────────────────────────────────
ROOT       = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DATA_DIR   = os.environ.get('WP6_DATA_DIR', os.path.join(ROOT, 'data', 'raw'))   # *_5min.duckdb files (see data/raw/README.md)
RESULTS    = os.path.join(ROOT, 'results')
FIGURES    = os.path.join(ROOT, 'figures')
os.makedirs(FIGURES, exist_ok=True)
os.makedirs(RESULTS, exist_ok=True)

plt.rcParams.update({
    'font.size': 10, 'axes.titlesize': 11, 'axes.labelsize': 10,
    'xtick.labelsize': 9, 'ytick.labelsize': 9, 'legend.fontsize': 9,
    'figure.dpi': 150, 'savefig.dpi': 150, 'font.family': 'DejaVu Sans',
    'axes.spines.top': False, 'axes.spines.right': False,
})

ASSET_COLORS = {
    'ES': '#1f77b4', 'NQ': '#aec7e8',
    'EURUSD': '#2ca02c', 'USDJPY': '#98df8a', 'GBPUSD': '#d62728',
    'GC': '#ff7f0e', 'CL': '#ffbb78',
    'BTC': '#9467bd', 'ETH': '#c5b0d5',
    'VIX': '#8c564b',
}
ASSET_CLASS = {
    'ES': 'Equity', 'NQ': 'Equity',
    'EURUSD': 'FX', 'USDJPY': 'FX', 'GBPUSD': 'FX',
    'GC': 'Commodity', 'CL': 'Commodity',
    'BTC': 'Crypto', 'ETH': 'Crypto',
    'VIX': 'Volatility',
}
CLASS_COLOR = {
    'Equity': '#1f77b4', 'FX': '#2ca02c',
    'Commodity': '#ff7f0e', 'Crypto': '#9467bd', 'Volatility': '#8c564b',
}

ASSETS = ['ES', 'NQ', 'EURUSD', 'USDJPY', 'GBPUSD', 'GC', 'CL', 'BTC', 'ETH', 'VIX']

# ─── 1. Load 5-min data ───────────────────────────────────────────────────────
print("=" * 70)
print("Loading 5-min OHLCV data...")

def load_asset(db_name, symbol):
    db_path = os.path.join(DATA_DIR, db_name)
    conn = duckdb.connect(db_path, read_only=True)
    df = conn.execute(
        f"SELECT datetime, close FROM ohlcv WHERE symbol='{symbol}' ORDER BY datetime"
    ).fetchdf()
    conn.close()
    df['datetime'] = pd.to_datetime(df['datetime'])
    df = df.set_index('datetime')
    df.columns = [symbol]
    return df

futures_assets = {
    'ES': 'futures_5min.duckdb',
    'NQ': 'futures_5min.duckdb',
    'GC': 'futures_5min.duckdb',
    'CL': 'futures_5min.duckdb',
    'VX': 'futures_5min.duckdb',
}
fx_assets = {
    'EURUSD': 'fx_5min.duckdb',
    'USDJPY': 'fx_5min.duckdb',
    'GBPUSD': 'fx_5min.duckdb',
}
crypto_assets = {
    'BTC': 'crypto_5min.duckdb',
    'ETH': 'crypto_5min.duckdb',
}
index_assets = {
    'VIX': 'index_5min.duckdb',
}

raw_data = {}
for sym, db in {**futures_assets, **fx_assets, **crypto_assets, **index_assets}.items():
    print(f"  Loading {sym} from {db}...")
    raw_data[sym] = load_asset(db, sym)

print("Data loaded.\n")

# ─── Compute 5-min log returns ────────────────────────────────────────────────
print("Computing 5-min log returns...")
log_returns = {}
for sym, df in raw_data.items():
    lr = np.log(df[sym] / df[sym].shift(1)).dropna()
    # Remove extreme outliers (> 20 std)
    mu, sigma = lr.mean(), lr.std()
    lr = lr[(lr - mu).abs() < 20 * sigma]
    log_returns[sym] = lr

# ─── 1. Descriptive Statistics ────────────────────────────────────────────────
print("Computing descriptive statistics...")
desc_rows = []
for sym in ASSET_CLASS.keys():
    if sym not in log_returns:
        # VX → treat VX as proxy for VIX
        continue
    lr = log_returns[sym]
    row = {
        'Asset': sym,
        'Class': ASSET_CLASS.get(sym, 'Other'),
        'N_obs': len(lr),
        'Date_Start': lr.index.min().strftime('%Y-%m-%d'),
        'Date_End': lr.index.max().strftime('%Y-%m-%d'),
        'Mean': lr.mean(),
        'Std': lr.std(),
        'Skewness': stats.skew(lr),
        'Kurtosis': stats.kurtosis(lr),
        'Min': lr.min(),
        'Max': lr.max(),
    }
    desc_rows.append(row)

desc_df = pd.DataFrame(desc_rows)
desc_df.to_csv(os.path.join(RESULTS, 'table1_descriptive_stats.csv'), index=False, float_format='%.6f')
print("  Saved table1_descriptive_stats.csv")
print(desc_df[['Asset', 'N_obs', 'Mean', 'Std', 'Skewness', 'Kurtosis']].to_string(index=False))

# ─── 2. Realized Volatility ───────────────────────────────────────────────────
print("\nComputing daily realized volatility...")

def compute_rv(log_ret_series, min_obs=20):
    """Annualized daily realized volatility."""
    df = log_ret_series.to_frame(name='ret')
    df['date'] = df.index.date
    grouped = df.groupby('date')['ret']
    rv = grouped.apply(lambda x: np.sqrt(np.sum(x**2)) * np.sqrt(252) if len(x) >= min_obs else np.nan)
    rv.index = pd.to_datetime(rv.index)
    return rv.dropna()

rv_dict = {}
for sym in ASSET_CLASS.keys():
    if sym not in log_returns:
        continue
    rv_dict[sym] = compute_rv(log_returns[sym])
    print(f"  {sym}: {len(rv_dict[sym])} trading days, RV mean = {rv_dict[sym].mean():.4f}")

# Combine into DataFrame
rv_df = pd.DataFrame(rv_dict)
rv_df.index.name = 'date'
rv_df.to_csv(os.path.join(RESULTS, 'daily_rv_all.csv'))
print("  Saved daily_rv_all.csv")

# ─── Fig 1: RV Time Series ────────────────────────────────────────────────────
print("\nPlotting Fig 1: RV Time Series...")
fig, axes = plt.subplots(5, 2, figsize=(14, 16), sharex=False)
axes = axes.flatten()

asset_order = ['ES', 'NQ', 'EURUSD', 'USDJPY', 'GBPUSD', 'GC', 'CL', 'BTC', 'ETH', 'VIX']
for i, sym in enumerate(asset_order):
    if sym not in rv_dict:
        continue
    ax = axes[i]
    rv_s = rv_dict[sym]
    color = ASSET_COLORS.get(sym, 'steelblue')
    ax.plot(rv_s.index, rv_s.values, color=color, linewidth=0.7, alpha=0.9)
    ax.fill_between(rv_s.index, rv_s.values, alpha=0.15, color=color)
    ax.set_title(f'{sym} ({ASSET_CLASS.get(sym, "")})')
    ax.set_ylabel('Annualized RV')
    ax.xaxis.set_major_locator(matplotlib.dates.YearLocator(2))
    ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter('%Y'))
    plt.setp(ax.xaxis.get_majorticklabels(), rotation=45, ha='right')

plt.suptitle('Daily Realized Volatility — All Assets (5-min, Annualized)', fontsize=13, y=1.01)
plt.tight_layout()
fig.savefig(os.path.join(FIGURES, 'fig1_rv_timeseries.png'), dpi=150, bbox_inches='tight')
plt.close()
print("  Saved fig1_rv_timeseries.png")

# ─── Helper: Diebold-Yilmaz GFEVD ────────────────────────────────────────────
def compute_gfevd(var_result, H=10):
    """
    Generalized Forecast Error Variance Decomposition (Pesaran-Shin).
    Returns an (n x n) matrix where entry [i,j] = share of forecast error
    variance of variable i attributable to shocks from variable j.
    """
    n = var_result.k_ar  # actually number of variables
    coefs = var_result.coefs         # shape (p, n, n)
    sigma = var_result.sigma_u       # (n, n) residual covariance
    n = sigma.shape[0]
    p = coefs.shape[0]

    # Build MA coefficient matrices Psi_h via recursive computation
    # Psi_0 = I, Psi_h = sum_{j=1..min(h,p)} A_j @ Psi_{h-j}
    psi = [np.eye(n)]
    for h in range(1, H + 1):
        psi_h = np.zeros((n, n))
        for j in range(1, min(h, p) + 1):
            psi_h += coefs[j-1] @ psi[h-j]
        psi.append(psi_h)

    # GFEVD: theta[i,j] = (1/sigma[j,j]) * sum_{h=0}^{H-1} (e_i' Psi_h Sigma e_j)^2
    #                      / sum_{h=0}^{H-1} (e_i' Psi_h Sigma Psi_h' e_i)
    sigma_diag = np.diag(sigma)
    gfevd = np.zeros((n, n))

    for i in range(n):
        ei = np.zeros(n); ei[i] = 1.0
        denom = 0.0
        for h in range(H):
            denom += ei @ psi[h] @ sigma @ psi[h].T @ ei
        for j in range(n):
            ej = np.zeros(n); ej[j] = 1.0
            num = 0.0
            for h in range(H):
                num += (ei @ psi[h] @ sigma @ ej) ** 2
            gfevd[i, j] = num / (sigma_diag[j] * denom)

    # Row-normalize so rows sum to 1
    row_sums = gfevd.sum(axis=1, keepdims=True)
    row_sums[row_sums == 0] = 1.0
    gfevd = gfevd / row_sums
    return gfevd


def compute_connectedness(rv_panel, assets, lag=4, H=10):
    """
    Given a DataFrame of RV series (rows=dates, cols=assets),
    fit VAR(lag), compute GFEVD at horizon H, return connectedness table.
    """
    # Use only complete rows
    data = rv_panel[assets].dropna()
    if len(data) < lag + H + 10:
        return None, None

    model = VAR(data)
    try:
        res = model.fit(lag, trend='c')
    except Exception as e:
        print(f"    VAR fit failed: {e}")
        return None, None

    gfevd = compute_gfevd(res, H=H)
    n = len(assets)

    # Build full connectedness table
    # Rows = "TO" (recipient of shocks), Cols = "FROM" (origin of shocks)
    # But DY convention: entry [i,j] = contribution of j to i's forecast error
    # TO_j = sum_{i != j} theta[i,j] / N * 100   (what j contributes to others)
    # FROM_i = sum_{j != i} theta[i,j] / N * 100  (what i receives from others)

    # gfevd[i,j] = share of i's FEVD from j's shock
    # So: FROM_i = sum_{j!=i} gfevd[i,j]  * 100
    #     TO_j   = sum_{i!=j} gfevd[i,j]  * 100

    # Total spillover index
    off_diag = gfevd.copy()
    np.fill_diagonal(off_diag, 0.0)
    TSI = off_diag.sum() / n * 100

    # Directional
    FROM = off_diag.sum(axis=1) * 100   # what each asset receives
    TO   = off_diag.sum(axis=0) * 100   # what each asset transmits
    NET  = TO - FROM

    # Build output DataFrame (DY table format)
    # gfevd * 100 → percentages
    tbl = pd.DataFrame(gfevd * 100, index=assets, columns=assets)
    tbl['FROM'] = FROM
    to_row = pd.Series(TO, index=assets)
    net_row = pd.Series(NET, index=assets)
    tbl.loc['TO']  = to_row
    tbl.loc['NET'] = net_row
    tbl.loc['TO', 'FROM']  = TSI
    tbl.loc['NET', 'FROM'] = np.nan

    return tbl, TSI


# ─── 3. Full-Sample Connectedness ─────────────────────────────────────────────
print("\n" + "="*70)
print("Section 3: Full-Sample Connectedness (VAR-4, H=10)...")

# Use VIX from index (not VX futures) as VIX is more common
# Find common sample: align all RV series
common_assets = [a for a in ASSET_CLASS.keys() if a in rv_dict]
rv_panel = pd.DataFrame({a: rv_dict[a] for a in common_assets})

# Forward fill up to 5 days to handle weekends/holidays
rv_panel = rv_panel.ffill(limit=5)

# Drop rows where more than 3 assets are missing
rv_panel = rv_panel.dropna(thresh=len(common_assets) - 2)
rv_panel = rv_panel.dropna()

print(f"  Common sample: {rv_panel.index[0].date()} to {rv_panel.index[-1].date()}, {len(rv_panel)} days")
print(f"  Assets: {common_assets}")

conn_table, tsi_full = compute_connectedness(rv_panel, common_assets, lag=4, H=10)
if conn_table is not None:
    conn_table.to_csv(os.path.join(RESULTS, 'table2_connectedness_full.csv'), float_format='%.2f')
    print(f"  Total Spillover Index (full sample): {tsi_full:.2f}%")
    print("  NET spillovers:")
    for a in common_assets:
        print(f"    {a:10s}: NET = {conn_table.loc['NET', a]:+.2f}%")
    print("  Saved table2_connectedness_full.csv")
else:
    print("  WARNING: Full-sample connectedness failed.")

# ─── Fig 2: Connectedness Heatmap ─────────────────────────────────────────────
if conn_table is not None:
    print("\nPlotting Fig 2: Connectedness Heatmap...")
    heat_data = conn_table.loc[common_assets, common_assets].values.astype(float)
    fig, ax = plt.subplots(figsize=(11, 8))
    mask_diag = np.eye(len(common_assets), dtype=bool)
    im = ax.imshow(heat_data, cmap='RdYlBu_r', aspect='auto', vmin=0, vmax=heat_data.max())
    for i in range(len(common_assets)):
        for j in range(len(common_assets)):
            txt = f"{heat_data[i,j]:.1f}"
            color = 'white' if heat_data[i,j] > heat_data.max() * 0.6 else 'black'
            ax.text(j, i, txt, ha='center', va='center', fontsize=8, color=color, fontweight='bold')
    ax.set_xticks(range(len(common_assets)))
    ax.set_yticks(range(len(common_assets)))
    ax.set_xticklabels(common_assets, rotation=45, ha='right')
    ax.set_yticklabels(common_assets)
    ax.set_xlabel('FROM (shock origin)')
    ax.set_ylabel('TO (shock recipient)')
    ax.set_title(f'Diebold-Yilmaz Connectedness Matrix (Full Sample)\nTotal Spillover Index = {tsi_full:.1f}%',
                 fontweight='bold')
    plt.colorbar(im, ax=ax, label='% of Forecast Error Variance')
    plt.tight_layout()
    fig.savefig(os.path.join(FIGURES, 'fig2_connectedness_heatmap.png'), dpi=150, bbox_inches='tight')
    plt.close()
    print("  Saved fig2_connectedness_heatmap.png")

# ─── 4. Rolling Connectedness ─────────────────────────────────────────────────
print("\n" + "="*70)
print("Section 4: Rolling Connectedness (200-day window, step=1)...")

WINDOW = 200
dates_all = rv_panel.index.tolist()
rolling_dates = []
rolling_tsi = []

n_steps = len(dates_all) - WINDOW
print(f"  Total windows to compute: {n_steps}")

for i in range(0, n_steps, 1):
    window_data = rv_panel.iloc[i : i + WINDOW]
    end_date = dates_all[i + WINDOW - 1]

    # Only compute if at least 150 non-NaN rows after dropna
    sub = window_data.dropna()
    if len(sub) < 100:
        continue

    model = VAR(sub)
    try:
        res = model.fit(4, trend='c')
        gfevd = compute_gfevd(res, H=10)
        n = gfevd.shape[0]
        off_diag = gfevd.copy()
        np.fill_diagonal(off_diag, 0.0)
        tsi = off_diag.sum() / n * 100
        rolling_dates.append(end_date)
        rolling_tsi.append(tsi)
    except Exception:
        pass

    if i % 200 == 0:
        print(f"    Step {i}/{n_steps}: {end_date.date()}, TSI = {rolling_tsi[-1]:.1f}%" if rolling_tsi else f"    Step {i}/{n_steps}")

rolling_df = pd.DataFrame({'date': rolling_dates, 'total_spillover': rolling_tsi})
rolling_df = rolling_df.set_index('date')
rolling_df.to_csv(os.path.join(RESULTS, 'rolling_total_spillover.csv'))
print(f"  Rolling analysis complete: {len(rolling_df)} observations")
print(f"  Mean TSI: {rolling_df['total_spillover'].mean():.2f}%")
print(f"  Min TSI:  {rolling_df['total_spillover'].min():.2f}%")
print(f"  Max TSI:  {rolling_df['total_spillover'].max():.2f}%")

# ─── Fig 3: Rolling Spillover ─────────────────────────────────────────────────
print("\nPlotting Fig 3: Rolling Total Spillover...")
fig, ax = plt.subplots(figsize=(13, 5))
ax.plot(rolling_df.index, rolling_df['total_spillover'], color='#1f77b4', linewidth=1.2, label='Total Spillover Index')
ax.fill_between(rolling_df.index, rolling_df['total_spillover'],
                rolling_df['total_spillover'].mean(), alpha=0.15, color='#1f77b4')
ax.axhline(rolling_df['total_spillover'].mean(), color='gray', linestyle='--', linewidth=0.8, label='Full-period mean')

events = {
    'COVID crash\n(Mar 2020)':  pd.Timestamp('2020-03-15'),
    'Fed hikes\n(Mar 2022)':    pd.Timestamp('2022-03-15'),
    'Crypto winter\n(Jun 2022)': pd.Timestamp('2022-06-01'),
    'FTX\n(Nov 2022)':          pd.Timestamp('2022-11-11'),
}
colors_ev = ['#d62728', '#ff7f0e', '#9467bd', '#8c564b']
for (label, date), c in zip(events.items(), colors_ev):
    if rolling_df.index.min() <= date <= rolling_df.index.max():
        ax.axvline(date, color=c, linestyle='--', linewidth=1.2, alpha=0.8)
        ax.text(date, ax.get_ylim()[1] if ax.get_ylim()[1] > 0 else rolling_df['total_spillover'].max(),
                label, ha='center', va='top', fontsize=7.5, color=c, rotation=0,
                bbox=dict(boxstyle='round,pad=0.2', facecolor='white', alpha=0.7, edgecolor=c, linewidth=0.5))

ax.set_xlabel('Date')
ax.set_ylabel('Total Spillover Index (%)')
ax.set_title('Rolling 200-Day Diebold-Yilmaz Total Spillover Index (VAR-4, H=10)', fontweight='bold')
ax.legend(loc='upper left')
ax.xaxis.set_major_locator(matplotlib.dates.YearLocator(1))
ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter('%Y'))
plt.setp(ax.xaxis.get_majorticklabels(), rotation=45, ha='right')
plt.tight_layout()
fig.savefig(os.path.join(FIGURES, 'fig3_rolling_spillover.png'), dpi=150, bbox_inches='tight')
plt.close()
print("  Saved fig3_rolling_spillover.png")

# ─── 5. Regime Analysis ───────────────────────────────────────────────────────
print("\n" + "="*70)
print("Section 5: Regime Analysis...")

regimes = {
    'Pre-COVID': ('2019-01-01', '2020-02-28', 'table3a_connectedness_precovid.csv'),
    'COVID crash': ('2020-03-01', '2020-06-30', 'table3b_connectedness_covid.csv'),
    'Rate hikes': ('2022-01-01', '2022-12-31', 'table3c_connectedness_ratehikes.csv'),
}

regime_tsi = {}
regime_net = {}

for regime_name, (start, end, fname) in regimes.items():
    sub = rv_panel.loc[start:end].dropna()
    print(f"\n  Regime: {regime_name} ({start} to {end}), {len(sub)} days")
    if len(sub) < 30:
        print("    WARNING: Too few observations, skipping.")
        continue
    # Use lag=2 if few obs
    lag = 4 if len(sub) >= 80 else 2
    tbl, tsi = compute_connectedness(sub, common_assets, lag=lag, H=10)
    if tbl is not None:
        tbl.to_csv(os.path.join(RESULTS, fname), float_format='%.2f')
        print(f"    TSI = {tsi:.2f}%")
        regime_tsi[regime_name] = tsi
        regime_net[regime_name] = {a: tbl.loc['NET', a] for a in common_assets}
        print(f"    NET spillovers: {regime_net[regime_name]}")
        print(f"    Saved {fname}")

# ─── Fig 4: NET spillovers by regime ─────────────────────────────────────────
print("\nPlotting Fig 4: NET spillovers by regime...")
if len(regime_net) > 0:
    fig, ax = plt.subplots(figsize=(12, 6))
    x = np.arange(len(common_assets))
    regime_list = list(regime_net.keys())
    n_regimes = len(regime_list)
    width = 0.8 / n_regimes
    regime_palette = ['#1f77b4', '#d62728', '#ff7f0e']

    for k, (rname, c) in enumerate(zip(regime_list, regime_palette)):
        nets = [regime_net[rname].get(a, np.nan) for a in common_assets]
        offset = (k - (n_regimes - 1) / 2) * width
        ax.bar(x + offset, nets, width=width * 0.9, label=rname, color=c, alpha=0.8, edgecolor='white', linewidth=0.5)

    ax.axhline(0, color='black', linewidth=0.8)
    ax.set_xticks(x)
    ax.set_xticklabels(common_assets)
    ax.set_xlabel('Asset')
    ax.set_ylabel('NET Spillover (%)')
    ax.set_title('NET Directional Spillovers by Market Regime', fontweight='bold')
    ax.legend(title='Regime')
    # Add TSI annotations
    for k, rname in enumerate(regime_list):
        if rname in regime_tsi:
            ax.text(0.02 + k * 0.18, 0.97, f'{rname}\nTSI={regime_tsi[rname]:.1f}%',
                    transform=ax.transAxes, fontsize=8, va='top', color=regime_palette[k],
                    bbox=dict(boxstyle='round', facecolor='white', alpha=0.7, edgecolor=regime_palette[k], linewidth=0.8))
    plt.tight_layout()
    fig.savefig(os.path.join(FIGURES, 'fig4_net_spillovers_regimes.png'), dpi=150, bbox_inches='tight')
    plt.close()
    print("  Saved fig4_net_spillovers_regimes.png")

# ─── 6. Trading Hour Analysis ─────────────────────────────────────────────────
print("\n" + "="*70)
print("Section 6: Trading Hour Analysis...")

session_assets = ['VIX', 'ES', 'NQ', 'BTC']

# Load 5-min data for session analysis
def load_full_asset(db_name, symbol):
    db_path = os.path.join(DATA_DIR, db_name)
    conn = duckdb.connect(db_path, read_only=True)
    df = conn.execute(
        f"SELECT datetime, close FROM ohlcv WHERE symbol='{symbol}' ORDER BY datetime"
    ).fetchdf()
    conn.close()
    df['datetime'] = pd.to_datetime(df['datetime'])
    df = df.set_index('datetime')
    # Ensure UTC-naive (data is already UTC)
    return df

session_data = {
    'ES':  load_full_asset('futures_5min.duckdb', 'ES'),
    'NQ':  load_full_asset('futures_5min.duckdb', 'NQ'),
    'VIX': load_full_asset('index_5min.duckdb', 'VIX'),
    'BTC': load_full_asset('crypto_5min.duckdb', 'BTC'),
}

def classify_session(hour):
    if 0 <= hour < 8:
        return 'Asian'
    elif 8 <= hour < 14.5:
        return 'European'
    elif 14.5 <= hour < 21:
        return 'US'
    else:
        return 'Other'

session_results = {}

for sym, df in session_data.items():
    df = df.copy()
    df['log_ret'] = np.log(df['close'] / df['close'].shift(1))
    df['abs_ret'] = df['log_ret'].abs()
    df['hour'] = df.index.hour + df.index.minute / 60
    df['session'] = df['hour'].apply(classify_session)
    df['hour_int'] = df.index.hour
    session_results[sym] = df

# Hourly average absolute returns
print("\n  Hourly average absolute returns by session:")
hourly_avgs = {}
for sym, df in session_results.items():
    hav = df.groupby('hour_int')['abs_ret'].mean() * 100  # in %
    hourly_avgs[sym] = hav

# Session-level statistics
session_stats = []
for sym, df in session_results.items():
    for sess in ['Asian', 'European', 'US']:
        sub = df[df['session'] == sess]
        if len(sub) > 10:
            session_stats.append({
                'Asset': sym,
                'Session': sess,
                'Mean_AbsRet': sub['abs_ret'].mean() * 100,
                'N_obs': len(sub),
            })
            print(f"    {sym:6s} | {sess:10s}: Mean |ret| = {sub['abs_ret'].mean()*100:.4f}%")

# Granger causality between VX/VIX, ES, NQ, BTC RV by session
print("\n  Granger causality tests by session (VAR-4)...")

# Use daily RV for granger tests (more robust)
gc_assets_rv = {a: rv_dict.get(a) for a in ['VIX', 'ES', 'NQ', 'BTC'] if a in rv_dict}
gc_panel = pd.DataFrame(gc_assets_rv).dropna()

granger_rows = []
pairs = [('VIX', 'ES'), ('VIX', 'NQ'), ('VIX', 'BTC'),
         ('ES', 'NQ'), ('ES', 'BTC'), ('NQ', 'BTC'),
         ('ES', 'VIX'), ('NQ', 'VIX'), ('BTC', 'VIX')]

# For session-specific granger, use 5-min abs returns aggregated to daily by session
for sess in ['Asian', 'European', 'US']:
    for (x_sym, y_sym) in pairs:
        if x_sym not in gc_assets_rv or y_sym not in gc_assets_rv:
            continue
        # Use full daily RV (session-specific daily RV would need more computation)
        try:
            pair_data = gc_panel[[y_sym, x_sym]].dropna()
            if len(pair_data) < 50:
                continue
            gc_test = grangercausalitytests(pair_data, maxlag=4, verbose=False)
            # F-stat from lag=4
            fstat = gc_test[4][0]['ssr_ftest'][0]
            pval  = gc_test[4][0]['ssr_ftest'][1]
            granger_rows.append({
                'Session': sess,
                'Cause': x_sym,
                'Effect': y_sym,
                'F_stat': fstat,
                'p_value': pval,
                'Significant': 'Yes' if pval < 0.05 else 'No',
            })
        except Exception as e:
            pass

granger_df = pd.DataFrame(granger_rows)
granger_df.to_csv(os.path.join(RESULTS, 'table4_session_granger.csv'), index=False, float_format='%.4f')
print("  Saved table4_session_granger.csv")

# ─── Fig 5: Hourly Volatility ─────────────────────────────────────────────────
print("\nPlotting Fig 5: Hourly Volatility Pattern...")
fig, axes = plt.subplots(2, 2, figsize=(13, 8))
axes = axes.flatten()
sym_list = ['VIX', 'ES', 'NQ', 'BTC']

for i, sym in enumerate(sym_list):
    ax = axes[i]
    hav = hourly_avgs[sym]
    color = ASSET_COLORS.get(sym, 'steelblue')

    bars = ax.bar(hav.index, hav.values, color=color, alpha=0.8, edgecolor='white', linewidth=0.5)

    # Session shading
    ax.axvspan(0, 8, alpha=0.07, color='blue', label='Asian (00-08 UTC)')
    ax.axvspan(8, 14.5, alpha=0.07, color='green', label='European (08-14:30 UTC)')
    ax.axvspan(14.5, 21, alpha=0.07, color='orange', label='US (14:30-21 UTC)')

    ax.set_title(f'{sym} — Avg |5-min return| by Hour (UTC)', fontweight='bold')
    ax.set_xlabel('Hour of Day (UTC)')
    ax.set_ylabel('Mean |Return| (%)')
    ax.set_xticks(range(0, 24, 2))
    if i == 0:
        ax.legend(fontsize=7.5, loc='upper right')

plt.suptitle('Intraday Volatility Pattern by Asset and Trading Session', fontsize=12, fontweight='bold')
plt.tight_layout()
fig.savefig(os.path.join(FIGURES, 'fig5_hourly_volatility.png'), dpi=150, bbox_inches='tight')
plt.close()
print("  Saved fig5_hourly_volatility.png")

# ─── 7. Net Directional Spillover Evolution ───────────────────────────────────
print("\n" + "="*70)
print("Section 7: Rolling NET spillovers for BTC and ES...")

rolling_net_dates = []
rolling_net_btc   = []
rolling_net_es    = []

n_steps2 = len(dates_all) - WINDOW
for i in range(0, n_steps2, 1):
    window_data = rv_panel.iloc[i : i + WINDOW]
    end_date = dates_all[i + WINDOW - 1]
    sub = window_data.dropna()
    if len(sub) < 100:
        continue

    model = VAR(sub)
    try:
        res = model.fit(4, trend='c')
        gfevd = compute_gfevd(res, H=10)
        n = gfevd.shape[0]
        off_diag = gfevd.copy()
        np.fill_diagonal(off_diag, 0.0)
        TO = off_diag.sum(axis=0) * 100
        FROM = off_diag.sum(axis=1) * 100
        NET = TO - FROM

        idx_btc = common_assets.index('BTC') if 'BTC' in common_assets else None
        idx_es  = common_assets.index('ES')  if 'ES'  in common_assets else None

        if idx_btc is not None and idx_es is not None:
            rolling_net_dates.append(end_date)
            rolling_net_btc.append(NET[idx_btc])
            rolling_net_es.append(NET[idx_es])
    except Exception:
        pass

    if i % 200 == 0 and rolling_net_btc:
        print(f"    Step {i}/{n_steps2}: BTC NET={rolling_net_btc[-1]:.1f}%, ES NET={rolling_net_es[-1]:.1f}%")

net_df = pd.DataFrame({
    'date': rolling_net_dates,
    'NET_BTC': rolling_net_btc,
    'NET_ES': rolling_net_es,
}).set_index('date')

# ─── Fig 6: Rolling NET BTC vs ES ─────────────────────────────────────────────
print("\nPlotting Fig 6: Rolling NET directional spillovers BTC vs ES...")
fig, ax = plt.subplots(figsize=(13, 5))
ax.plot(net_df.index, net_df['NET_BTC'], color=ASSET_COLORS['BTC'], linewidth=1.1, label='BTC NET spillover', alpha=0.9)
ax.plot(net_df.index, net_df['NET_ES'],  color=ASSET_COLORS['ES'],  linewidth=1.1, label='ES NET spillover',  alpha=0.9)
ax.axhline(0, color='black', linewidth=0.8, linestyle='--')
ax.fill_between(net_df.index, net_df['NET_BTC'], 0,
                where=(net_df['NET_BTC'] > 0), alpha=0.15, color=ASSET_COLORS['BTC'], label='BTC net transmitter')
ax.fill_between(net_df.index, net_df['NET_BTC'], 0,
                where=(net_df['NET_BTC'] < 0), alpha=0.15, color='gray', label='BTC net receiver')

for (label, date), c in zip(events.items(), colors_ev):
    if net_df.index.min() <= date <= net_df.index.max():
        ax.axvline(date, color=c, linestyle=':', linewidth=1, alpha=0.8)
        ax.text(date, ax.get_ylim()[1] if ax.get_ylim()[1] else net_df[['NET_BTC','NET_ES']].max().max(),
                label, ha='center', va='top', fontsize=7, color=c,
                bbox=dict(boxstyle='round,pad=0.2', facecolor='white', alpha=0.6))

ax.set_xlabel('Date')
ax.set_ylabel('NET Directional Spillover (%)')
ax.set_title('Rolling 200-Day NET Directional Spillovers: BTC vs ES (VAR-4, H=10)', fontweight='bold')
ax.legend(loc='lower right', ncol=2)
ax.xaxis.set_major_locator(matplotlib.dates.YearLocator(1))
ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter('%Y'))
plt.setp(ax.xaxis.get_majorticklabels(), rotation=45, ha='right')
plt.tight_layout()
fig.savefig(os.path.join(FIGURES, 'fig6_net_directional_btc_es.png'), dpi=150, bbox_inches='tight')
plt.close()
print("  Saved fig6_net_directional_btc_es.png")

# ─── 8. Network Visualization ─────────────────────────────────────────────────
print("\n" + "="*70)
print("Section 8: Network Graph...")

if conn_table is not None:
    gfevd_mat = conn_table.loc[common_assets, common_assets].values.astype(float)
    np.fill_diagonal(gfevd_mat, 0.0)

    # Threshold: keep edges above median
    all_edges = gfevd_mat.flatten()
    all_edges = all_edges[all_edges > 0]
    threshold = np.median(all_edges)

    G = nx.DiGraph()
    for a in common_assets:
        G.add_node(a, asset_class=ASSET_CLASS.get(a, 'Other'))

    TO_vals = {}
    for j, asj in enumerate(common_assets):
        TO_vals[asj] = sum(gfevd_mat[i, j] for i in range(len(common_assets)) if i != j)

    for i, asi in enumerate(common_assets):
        for j, asj in enumerate(common_assets):
            if i != j and gfevd_mat[i, j] > threshold:
                G.add_edge(asj, asi, weight=gfevd_mat[i, j])  # j → i (j causes forecast error in i)

    # Layout
    pos = nx.spring_layout(G, k=2.5, seed=42, iterations=100)

    node_sizes   = [max(300, TO_vals.get(a, 1) * 40) for a in G.nodes()]
    node_colors  = [CLASS_COLOR.get(ASSET_CLASS.get(a, 'Other'), 'gray') for a in G.nodes()]
    edge_weights = [G[u][v]['weight'] for u, v in G.edges()]
    max_w = max(edge_weights) if edge_weights else 1
    edge_widths  = [w / max_w * 4 for w in edge_weights]
    edge_alphas  = [0.3 + 0.5 * w / max_w for w in edge_weights]

    fig, ax = plt.subplots(figsize=(12, 9))

    for (u, v), lw, alpha, w in zip(G.edges(), edge_widths, edge_alphas, edge_weights):
        nx.draw_networkx_edges(G, pos, edgelist=[(u, v)], ax=ax,
                               width=lw, alpha=alpha, edge_color='#555555',
                               arrows=True, arrowsize=12,
                               connectionstyle='arc3,rad=0.15',
                               node_size=node_sizes[list(G.nodes()).index(u)] if u in G.nodes() else 300)

    nx.draw_networkx_nodes(G, pos, ax=ax,
                           node_size=node_sizes,
                           node_color=node_colors,
                           alpha=0.9, edgecolors='white', linewidths=1.5)
    nx.draw_networkx_labels(G, pos, ax=ax, font_size=9, font_weight='bold', font_color='black')

    # Legend for asset classes
    legend_patches = [mpatches.Patch(color=c, label=cls)
                      for cls, c in CLASS_COLOR.items()]
    ax.legend(handles=legend_patches, title='Asset Class', loc='upper left', fontsize=9)
    ax.set_title(f'Volatility Transmission Network\n(Edges > median threshold = {threshold:.1f}%,\n'
                 f'Node size ∝ Total TO spillover, Arrow direction = shock transmission)',
                 fontweight='bold', fontsize=10)
    ax.axis('off')
    plt.tight_layout()
    fig.savefig(os.path.join(FIGURES, 'fig7_network_graph.png'), dpi=150, bbox_inches='tight')
    plt.close()
    print("  Saved fig7_network_graph.png")

# ─── Summary Report ───────────────────────────────────────────────────────────
print("\n" + "=" * 70)
print("ANALYSIS COMPLETE — KEY NUMERICAL FINDINGS")
print("=" * 70)

print("\n[Table 1] Descriptive Statistics (5-min log returns):")
for row in desc_rows:
    print(f"  {row['Asset']:10s}: N={row['N_obs']:>8,}, Std={row['Std']:.6f}, "
          f"Skew={row['Skewness']:+.3f}, Kurt={row['Kurtosis']:+.2f}")

print(f"\n[Section 2] Daily RV — Annualized means:")
for sym, rv in rv_dict.items():
    print(f"  {sym:10s}: mean RV = {rv.mean():.4f}, max RV = {rv.max():.4f}")

print(f"\n[Section 3] Full-Sample Connectedness:")
print(f"  Total Spillover Index = {tsi_full:.2f}%" if tsi_full else "  N/A")
if conn_table is not None:
    print("  NET spillovers (positive = net transmitter):")
    for a in common_assets:
        net_val = conn_table.loc['NET', a]
        to_val  = conn_table.loc['TO', a]
        frm_val = conn_table.loc[a, 'FROM']
        print(f"    {a:10s}: TO={to_val:+.2f}%, FROM={frm_val:+.2f}%, NET={net_val:+.2f}%")

print(f"\n[Section 4] Rolling Connectedness ({len(rolling_df)} windows):")
print(f"  Mean TSI: {rolling_df['total_spillover'].mean():.2f}%")
print(f"  Max TSI:  {rolling_df['total_spillover'].max():.2f}% on {rolling_df['total_spillover'].idxmax().date()}")
print(f"  Min TSI:  {rolling_df['total_spillover'].min():.2f}% on {rolling_df['total_spillover'].idxmin().date()}")

print(f"\n[Section 5] Regime TSI:")
for r, tsi in regime_tsi.items():
    print(f"  {r:20s}: TSI = {tsi:.2f}%")

print(f"\n[Section 6] Granger causality (sample):")
sig_gc = granger_df[granger_df['Significant'] == 'Yes'] if len(granger_df) > 0 else pd.DataFrame()
print(f"  {len(sig_gc)} / {len(granger_df)} pairs significant at 5%")
if len(sig_gc) > 0:
    print(sig_gc[['Session', 'Cause', 'Effect', 'F_stat', 'p_value']].head(10).to_string(index=False))

print(f"\n[Section 7] Rolling NET BTC (last window): {rolling_net_btc[-1]:.2f}%" if rolling_net_btc else "")
print(f"  Rolling NET ES  (last window): {rolling_net_es[-1]:.2f}%" if rolling_net_es else "")

print("\nAll outputs saved to:", RESULTS)
print("=" * 70)