!pip install yfinance pandas numpy matplotlib seaborn scipy --quiet

import yfinance as yf
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import stats
import warnings

warnings.filterwarnings('ignore')
%matplotlib inline
plt.rcParams['figure.figsize'] = (14, 5)
plt.rcParams['axes.spines.top']   = False
plt.rcParams['axes.spines.right'] = False
sns.set_palette('deep')
print('Imports ready.')

START_DATE      = '2017-01-01'
ROLLING_WINDOWS = [30, 60, 90]
TICKERS = {
    'SPX':    '^GSPC',
    'NASDAQ': '^IXIC',
    'VIX':    '^VIX',
    'BTC':    'BTC-USD',
    'ETH':    'ETH-USD',
}
USE_SYNTHETIC = False

MACRO_EVENTS = [
    {'date': '2020-03-16', 'label': 'COVID Crash Low'},
    {'date': '2021-11-22', 'label': 'Market Peak'},
    {'date': '2022-01-26', 'label': 'Fed Pivot Signal'},
    {'date': '2022-06-15', 'label': '75bps Hike'},
    {'date': '2023-01-13', 'label': 'Bear Mkt Bottom'},
    {'date': '2024-01-10', 'label': 'BTC ETF Approved'},
]

def fetch_multi_asset(
    tickers: dict,
    start: str
) -> pd.DataFrame:
    """
    Fetch daily close prices for multiple assets and align them.

    Parameters
    ----------
    tickers : dict
        Mapping of friendly names to Yahoo Finance tickers.
        Example: {'SPX': '^GSPC', 'BTC': 'BTC-USD'}.
    start : str
        Start date in 'YYYY-MM-DD' format.

    Returns
    -------
    pd.DataFrame
        Daily close prices with friendly column names.
        VIX is kept as its own column (levels, not returns used for conditioning).

    Notes
    -----
    Equity indices only trade on business days while BTC trades 24/7.
    Equity columns are forward-filled on weekends for alignment.
    This creates a slight look-ahead bias on weekend crypto moves — be aware
    when computing same-day correlations.
    """
    symbols = list(tickers.values())
    raw = yf.download(symbols, start=start, progress=False, auto_adjust=True)
    prices = raw['Close'].copy()
    reverse = {v: k for k, v in tickers.items()}
    prices.rename(columns=reverse, inplace=True)
    prices.index = pd.to_datetime(prices.index)
    prices = prices.ffill().dropna()
    print(f'Fetched {len(prices)} rows for {list(prices.columns)}')
    return prices


def generate_synthetic_multi_asset(start: str, n_days: int = 2000) -> pd.DataFrame:
    """
    Generate synthetic multi-asset prices with regime-varying correlation.

    Parameters
    ----------
    start : str
        Start date in 'YYYY-MM-DD' format.
    n_days : int
        Number of business days to simulate.

    Returns
    -------
    pd.DataFrame
        Synthetic prices for SPX, NASDAQ, VIX, BTC, ETH.
    """
    np.random.seed(7)
    dates = pd.date_range(start, periods=n_days, freq='B')

    # SPX and NASDAQ highly correlated (baseline ~0.92)
    spx_rets  = 0.0003 + 0.01 * np.random.randn(n_days)
    ixic_rets = 0.0004 + 0.012 * (0.92 * (spx_rets / 0.01) + 0.39 * np.random.randn(n_days)) * 0.01

    # BTC: regime-varying correlation with SPX
    corr_regime = np.concatenate([
        np.full(400, 0.15),
        np.full(400, 0.70),
        np.full(400, 0.80),
        np.full(400, 0.45),
        np.full(400, 0.30),
    ])[:n_days]

    spx_std = 0.01
    btc_rets = np.array([
        0.0005 + 0.035 * (
            corr_regime[i] * (spx_rets[i] / spx_std) + np.sqrt(1 - corr_regime[i]**2) * np.random.randn()
        )
        for i in range(n_days)
    ])
    eth_rets = btc_rets * 1.2 + 0.005 * np.random.randn(n_days)

    # VIX (inverse relationship with SPX returns roughly)
    vix_levels = 15 + 10 * np.clip(-np.cumsum(spx_rets) * 5 + np.random.randn(n_days) * 2, -10, 60)

    df = pd.DataFrame({
        'SPX':    2500 * np.exp(np.cumsum(spx_rets)),
        'NASDAQ': 6000 * np.exp(np.cumsum(ixic_rets)),
        'VIX':    np.abs(vix_levels),
        'BTC':    10000 * np.exp(np.cumsum(btc_rets)),
        'ETH':    300 * np.exp(np.cumsum(eth_rets)),
    }, index=dates)
    return df


if USE_SYNTHETIC:
    prices = generate_synthetic_multi_asset(START_DATE)
    print('Using synthetic data.')
else:
    try:
        prices = fetch_multi_asset(TICKERS, START_DATE)
    except Exception as e:
        print(f'Live fetch failed ({e}). Falling back to synthetic.')
        prices = generate_synthetic_multi_asset(START_DATE)

print(prices.tail(3))

def compute_cross_asset_metrics(
    prices: pd.DataFrame,
    windows: list,
    target: str = 'BTC',
    benchmarks: list = None
) -> pd.DataFrame:
    """
    Compute rolling correlation and beta between target asset and benchmarks.

    Parameters
    ----------
    prices : pd.DataFrame
        Multi-asset daily close prices.
    windows : list of int
        Rolling window sizes in days.
    target : str
        Asset whose correlation/beta we are measuring (e.g., 'BTC').
    benchmarks : list of str, optional
        Benchmark assets (e.g., ['SPX', 'NASDAQ']). Defaults to all except target.

    Returns
    -------
    pd.DataFrame
        Rolling correlations and betas indexed by date.
        Beta = rolling_cov(target, benchmark) / rolling_var(benchmark).

    Notes
    -----
    Beta > 1 means BTC amplifies equity moves.
    Beta < 0 means BTC moves opposite to equities (rare but occurs during crypto bull markets).
    Rolling beta is noisy on short windows; 60d is a reasonable trade-off.
    """
    if benchmarks is None:
        benchmarks = [c for c in prices.columns if c not in [target, 'VIX']]

    log_rets = np.log(prices / prices.shift(1)).dropna()
    metrics = pd.DataFrame(index=log_rets.index)

    for bm in benchmarks:
        for w in windows:
            metrics[f'corr_{bm}_{w}d'] = log_rets[target].rolling(w).corr(log_rets[bm])

        # Rolling beta at 60d
        w = 60
        cov = log_rets[target].rolling(w).cov(log_rets[bm])
        var = log_rets[bm].rolling(w).var()
        metrics[f'beta_{bm}_60d'] = cov / var

    return metrics


metrics = compute_cross_asset_metrics(prices, ROLLING_WINDOWS)
print('Cross-asset metrics computed.')
print(metrics[['corr_SPX_60d', 'corr_NASDAQ_60d', 'beta_SPX_60d']].tail(5))

def compute_lead_lag(
    prices: pd.DataFrame,
    asset_a: str = 'SPX',
    asset_b: str = 'BTC',
    max_lag: int = 5
) -> pd.DataFrame:
    """
    Compute cross-correlation at various lags to identify lead-lag relationship.

    Parameters
    ----------
    prices : pd.DataFrame
        Multi-asset close prices.
    asset_a : str
        First asset name (candidate leader).
    asset_b : str
        Second asset name (candidate follower).
    max_lag : int
        Maximum lag in days (both positive and negative).

    Returns
    -------
    pd.DataFrame
        Cross-correlation at each lag. Positive lag k means A leads B by k days.

    Notes
    -----
    A positive correlation at lag +k means asset_a's return today predicts
    asset_b's return k days from now — asset_a leads.
    Interpret with caution: most liquid markets adjust quickly and lead-lag
    relationships often disappear after transaction costs.
    """
    log_rets = np.log(prices / prices.shift(1)).dropna()
    a_rets = log_rets[asset_a]
    b_rets = log_rets[asset_b]

    results = []
    for lag in range(-max_lag, max_lag + 1):
        if lag >= 0:
            corr = a_rets.corr(b_rets.shift(-lag))
            label = f'{asset_a} leads {asset_b} by {lag}d'
        else:
            corr = b_rets.corr(a_rets.shift(lag))
            label = f'{asset_b} leads {asset_a} by {abs(lag)}d'
        results.append({'lag': lag, 'correlation': round(corr, 4), 'interpretation': label})

    return pd.DataFrame(results)


lead_lag_df = compute_lead_lag(prices, 'SPX', 'BTC')
print('Lead-lag analysis (SPX vs BTC):')
print(lead_lag_df.to_string(index=False))

def plot_correlation_timeline(
    prices: pd.DataFrame,
    metrics: pd.DataFrame,
    macro_events: list
) -> None:
    """
    Plot normalized price history and rolling correlation timelines.

    Parameters
    ----------
    prices : pd.DataFrame
        Multi-asset close prices.
    metrics : pd.DataFrame
        Rolling correlation and beta output.
    macro_events : list of dict
        Each dict: {'date': 'YYYY-MM-DD', 'label': 'Event Name'}.
    """
    fig, axes = plt.subplots(2, 1, figsize=(15, 10), sharex=True)

    # Panel 1: Prices
    norm = prices[['SPX', 'BTC']] / prices[['SPX', 'BTC']].iloc[0] * 100
    axes[0].plot(norm.index, norm['SPX'], label='S&P 500 (rebased)', color='steelblue', linewidth=1.5)
    axes[0].plot(norm.index, norm['BTC'], label='BTC (rebased)',      color='orange',   linewidth=1.5)
    axes[0].set_yscale('log')
    axes[0].set_ylabel('Rebased Price (log scale)')
    axes[0].set_title('S&P 500 vs BTC — Performance Comparison (Base = 100)')
    axes[0].legend()

    # Panel 2: Rolling correlations
    for col, color in [('corr_SPX_30d', 'lightblue'), ('corr_SPX_60d', 'steelblue'), ('corr_SPX_90d', 'navy')]:
        axes[1].plot(metrics.index, metrics[col], label=col, color=color, linewidth=1.2, alpha=0.85)
    axes[1].axhline(0, color='black', linewidth=0.8, linestyle='--')
    axes[1].axhline(0.5, color='green', linewidth=0.5, linestyle=':', alpha=0.7)
    axes[1].set_ylim(-1, 1)
    axes[1].set_ylabel('Pearson Correlation')
    axes[1].set_title('Rolling BTC-SPX Log-Return Correlation')
    axes[1].legend()

    for event in macro_events:
        edate = pd.to_datetime(event['date'])
        for ax in axes:
            ax.axvline(edate, color='crimson', alpha=0.4, linewidth=1.0, linestyle='--')

    plt.tight_layout()
    plt.show()


def plot_correlation_by_vix_regime(
    prices: pd.DataFrame,
    metrics: pd.DataFrame
) -> None:
    """
    Compare BTC-SPX correlation in low, medium, and high VIX regimes.

    Parameters
    ----------
    prices : pd.DataFrame
        Must include 'VIX' column.
    metrics : pd.DataFrame
        Must include 'corr_SPX_60d' column.
    """
    combined = metrics[['corr_SPX_60d']].join(prices['VIX'], how='inner').dropna()
    vix_q = combined['VIX'].quantile([0.33, 0.67]).values
    combined['vix_regime'] = pd.cut(
        combined['VIX'],
        bins=[-np.inf, vix_q[0], vix_q[1], np.inf],
        labels=['Low Fear', 'Medium Fear', 'High Fear']
    )

    fig, ax = plt.subplots(figsize=(10, 5))
    colors = ['green', 'goldenrod', 'red']
    sns.boxplot(data=combined, x='vix_regime', y='corr_SPX_60d',
                order=['Low Fear', 'Medium Fear', 'High Fear'],
                palette=colors, ax=ax)
    ax.axhline(0, color='black', linewidth=0.8, linestyle='--')
    ax.set_title('BTC-SPX Correlation by VIX Fear Regime')
    ax.set_xlabel('VIX Regime')
    ax.set_ylabel('60d Rolling Correlation')
    plt.tight_layout()
    plt.show()

    print('\nMean correlation by VIX regime:')
    print(combined.groupby('vix_regime')['corr_SPX_60d'].agg(['mean', 'std', 'count']))


def plot_lead_lag(
    lead_lag_df: pd.DataFrame,
    asset_a: str = 'SPX',
    asset_b: str = 'BTC'
) -> None:
    """
    Bar chart of lead-lag cross-correlations.

    Parameters
    ----------
    lead_lag_df : pd.DataFrame
        Output of compute_lead_lag().
    asset_a, asset_b : str
        Asset names for labeling.
    """
    fig, ax = plt.subplots(figsize=(10, 4))
    colors = ['red' if x < 0 else 'steelblue' for x in lead_lag_df['correlation']]
    ax.bar(lead_lag_df['lag'], lead_lag_df['correlation'], color=colors, edgecolor='white')
    ax.axhline(0, color='black', linewidth=0.8)
    ax.axvline(0, color='black', linewidth=0.8, linestyle='--')
    ax.set_xlabel(f'Lag (days) — Positive = {asset_a} leads {asset_b}')
    ax.set_ylabel('Cross-Correlation')
    ax.set_title(f'{asset_a} vs {asset_b} Lead-Lag Analysis')
    plt.tight_layout()
    plt.show()


plot_correlation_timeline(prices, metrics, MACRO_EVENTS)
plot_correlation_by_vix_regime(prices, metrics)
plot_lead_lag(lead_lag_df)

def build_equity_crypto_regime_signal(
    metrics: pd.DataFrame,
    prices: pd.DataFrame,
    high_corr_threshold: float = 0.55,
    vix_risk_off_threshold: float = 30.0
) -> pd.DataFrame:
    """
    Build a composite regime signal combining SPX correlation and VIX.

    Parameters
    ----------
    metrics : pd.DataFrame
        Must include 'corr_SPX_60d' and 'beta_SPX_60d'.
    prices : pd.DataFrame
        Must include 'VIX' column.
    high_corr_threshold : float
        60d correlation above this = 'risk_on_correlated' regime.
    vix_risk_off_threshold : float
        VIX above this level = elevated systemic risk.

    Returns
    -------
    pd.DataFrame
        Columns: corr_60d, beta_60d, vix, regime.
        regime: 'correlated_high_fear', 'correlated_low_fear', 'decoupled'.

    Notes
    -----
    'correlated_high_fear' = BTC moving with equities during fear — dangerous zone.
    'correlated_low_fear'  = BTC moving with equities during calm — typical bull run.
    'decoupled'            = BTC on its own cycle — crypto-specific narrative driving price.
    """
    df = metrics[['corr_SPX_60d', 'beta_SPX_60d']].join(prices['VIX'], how='inner').dropna()
    df.columns = ['corr_60d', 'beta_60d', 'vix']

    conditions = [
        (df['corr_60d'] >= high_corr_threshold) & (df['vix'] >= vix_risk_off_threshold),
        (df['corr_60d'] >= high_corr_threshold) & (df['vix'] <  vix_risk_off_threshold),
    ]
    choices = ['correlated_high_fear', 'correlated_low_fear']
    df['regime'] = np.select(conditions, choices, default='decoupled')

    print('Regime distribution:')
    print(df['regime'].value_counts())
    return df


regime_signal = build_equity_crypto_regime_signal(metrics, prices)
print(regime_signal.tail(5))

def export_equity_crypto_data(
    prices: pd.DataFrame,
    metrics: pd.DataFrame,
    regime_signal: pd.DataFrame,
    lead_lag_df: pd.DataFrame
) -> None:
    """
    Export all computed data to CSV files.

    Parameters
    ----------
    prices, metrics, regime_signal, lead_lag_df : pd.DataFrame
        Data to export.
    """
    prices.to_csv('equity_crypto_prices.csv')
    metrics.to_csv('equity_crypto_correlations.csv')
    regime_signal.to_csv('equity_crypto_regime_signal.csv')
    lead_lag_df.to_csv('spx_btc_lead_lag.csv', index=False)
    print('Exported: equity_crypto_prices.csv')
    print('Exported: equity_crypto_correlations.csv')
    print('Exported: equity_crypto_regime_signal.csv')
    print('Exported: spx_btc_lead_lag.csv')


export_equity_crypto_data(prices, metrics, regime_signal, lead_lag_df)