!pip install numpy pandas matplotlib seaborn scipy --quiet

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

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

BUCKET_SIZE        = 50     # total trades per VPIN volume bucket
VPIN_WINDOW        = 10     # number of buckets to average for VPIN
VPIN_HIGH_THRESH   = 0.45   # VPIN above this → high toxicity
VPIN_LOW_THRESH    = 0.20   # VPIN below this → low toxicity

NORMAL_SPREAD_BPS  = 8.0
HIGH_TOX_SPREAD    = 25.0
BASE_QUOTE_SIZE    = 0.10

SIMULATION_TRADES  = 10_000
START_PRICE        = 50_000.0

def generate_synthetic_trades(
    n_trades: int,
    start_price: float
) -> pd.DataFrame:
    """
    Generate synthetic trade data with embedded high-toxicity episodes.

    During toxicity episodes, order flow becomes heavily one-sided and
    the mid-price moves adversely. This mimics whale accumulation/distribution
    or liquidation cascade scenarios.

    Parameters
    ----------
    n_trades : int
        Total number of trades to generate.
    start_price : float
        Initial mid-price.

    Returns
    -------
    pd.DataFrame
        Trade-level data: trade_id, price, side, size, mid_price, is_toxic_episode.

    Notes
    -----
    The synthetic data includes three embedded toxicity events where buy or sell
    pressure dominates for a sustained period. These events are 'labeled' in
    the output for validation purposes — in real data, we don't know in advance
    which trades are from informed sources.
    """
    np.random.seed(42)

    # Define toxicity event windows (trade range, direction, price impact)
    toxic_events = [
        (2000, 2300, 'buy',  0.03),   # informed buying → price up 3%
        (5000, 5200, 'sell', -0.04),  # aggressive selling → price down 4%
        (8000, 8250, 'sell', -0.05),  # large liquidation event
    ]

    mid    = start_price
    trades = []
    tick_vol = 0.02 / np.sqrt(5000)  # per-trade vol

    for i in range(n_trades):
        # Check if in a toxic episode
        is_toxic = False
        toxic_bias = 0.0
        for t_start, t_end, direction, price_impact in toxic_events:
            if t_start <= i < t_end:
                is_toxic = True
                toxic_bias = 0.85 if direction == 'buy' else 0.15  # buy probability
                # Gradually move price in the direction
                mid *= (1 + price_impact / (t_end - t_start))
                break

        # Normal price walk
        if not is_toxic:
            mid    *= np.exp(np.random.normal(0, tick_vol))
            buy_prob = 0.50
        else:
            buy_prob = toxic_bias

        mid = max(mid, 1.0)

        side = 'buy' if np.random.rand() < buy_prob else 'sell'
        size = np.abs(np.random.normal(0.05, 0.02))  # BTC size per trade

        # Trade price slightly offset from mid
        price = mid * (1.0001 if side == 'sell' else 0.9999)  # buys hit ask, sells hit bid

        trades.append({
            'trade_id':        i,
            'price':           mid,
            'side':            side,
            'size':            max(0.001, round(size, 3)),
            'mid_price':       mid,
            'is_toxic_episode': int(is_toxic),
        })

    df = pd.DataFrame(trades)
    print(f'Generated {len(df)} trades. Side distribution: {df["side"].value_counts().to_dict()}')
    print(f'Toxic episodes cover {df["is_toxic_episode"].sum()} trades ({df["is_toxic_episode"].mean()*100:.1f}%)')
    return df


trades = generate_synthetic_trades(SIMULATION_TRADES, START_PRICE)

def compute_vpin(
    trades: pd.DataFrame,
    bucket_size: int,
    vpin_window: int
) -> pd.DataFrame:
    """
    Compute VPIN (Volume-Synchronized Probability of Informed Trading).

    Groups trades into volume buckets of fixed size, then computes the
    volume imbalance within each bucket. VPIN is the rolling average
    of absolute imbalances across vpin_window buckets.

    Parameters
    ----------
    trades : pd.DataFrame
        Trade data with 'side' and 'size' columns.
    bucket_size : int
        Number of trades per volume bucket.
    vpin_window : int
        Number of buckets to average over for VPIN estimate.

    Returns
    -------
    pd.DataFrame
        Per-bucket DataFrame with: bucket_id, total_volume, buy_volume, sell_volume,
        imbalance, vpin, last_trade_id, avg_price, toxic_episode_flag.

    Notes
    -----
    The original VPIN paper uses dollar-volume buckets to normalize for price changes
    over time. This simplified version uses trade-count buckets, which is sufficient
    for educational purposes and works well for a single asset over short windows.
    """
    buckets = []
    n_trades = len(trades)

    for bucket_start in range(0, n_trades, bucket_size):
        bucket = trades.iloc[bucket_start:bucket_start + bucket_size]
        if len(bucket) < bucket_size // 2:
            continue

        buys  = bucket[bucket['side'] == 'buy']
        sells = bucket[bucket['side'] == 'sell']

        buy_vol   = buys['size'].sum()
        sell_vol  = sells['size'].sum()
        total_vol = buy_vol + sell_vol
        imbalance = abs(buy_vol - sell_vol) / (total_vol + 1e-8)

        buckets.append({
            'bucket_id':          len(buckets),
            'total_volume':       round(total_vol, 4),
            'buy_volume':         round(buy_vol, 4),
            'sell_volume':        round(sell_vol, 4),
            'imbalance':          round(imbalance, 4),
            'last_trade_id':      bucket.index[-1],
            'avg_price':          round(bucket['mid_price'].mean(), 2),
            'toxic_episode_flag': int(bucket['is_toxic_episode'].mean() > 0.5),
        })

    df = pd.DataFrame(buckets)
    df['vpin'] = df['imbalance'].rolling(vpin_window).mean()
    df['vpin_regime'] = pd.cut(
        df['vpin'].fillna(0),
        bins=[-0.01, VPIN_LOW_THRESH, VPIN_HIGH_THRESH, 1.01],
        labels=['low', 'medium', 'high']
    )

    print(f'VPIN computed over {len(df)} buckets.')
    print(f'Mean VPIN: {df["vpin"].mean():.3f} | Max: {df["vpin"].max():.3f}')
    print(f'Regime distribution:')
    print(df['vpin_regime'].value_counts().to_string())
    return df


vpin_df = compute_vpin(trades, BUCKET_SIZE, VPIN_WINDOW)
print(vpin_df[['bucket_id', 'imbalance', 'vpin', 'vpin_regime', 'toxic_episode_flag']].tail(10))

def validate_vpin_predictiveness(
    vpin_df: pd.DataFrame,
    lookahead_buckets: int = 3
) -> pd.DataFrame:
    """
    Test whether high VPIN buckets predict adverse price moves.

    Parameters
    ----------
    vpin_df : pd.DataFrame
        Output of compute_vpin().
    lookahead_buckets : int
        Number of buckets ahead to measure price change.

    Returns
    -------
    pd.DataFrame
        Summary of future price moves by VPIN regime.
    """
    df = vpin_df.copy()
    df['future_price_change'] = df['avg_price'].pct_change(lookahead_buckets).shift(-lookahead_buckets) * 100
    df['abs_future_change']   = df['future_price_change'].abs()

    summary = df.groupby('vpin_regime', observed=True).agg(
        n_buckets      = ('bucket_id', 'count'),
        mean_abs_move  = ('abs_future_change', 'mean'),
        max_abs_move   = ('abs_future_change', 'max'),
        toxic_rate     = ('toxic_episode_flag', 'mean'),
    ).round(3)

    print(f'\nVPIN Predictiveness ({lookahead_buckets}-bucket ahead price change):')
    print(summary.to_string())
    return summary


def compute_mm_toxicity_response(
    vpin_df: pd.DataFrame
) -> pd.DataFrame:
    """
    Map VPIN regime to market maker quote parameters.

    Parameters
    ----------
    vpin_df : pd.DataFrame
        VPIN output with 'vpin_regime' column.

    Returns
    -------
    pd.DataFrame
        vpin_df with added spread_bps, size_bps columns.
    """
    regime_params = {
        'low':    {'spread_bps': NORMAL_SPREAD_BPS, 'size_frac': 1.00},
        'medium': {'spread_bps': NORMAL_SPREAD_BPS * 1.5, 'size_frac': 0.70},
        'high':   {'spread_bps': HIGH_TOX_SPREAD,  'size_frac': 0.30},
    }

    out = vpin_df.copy()
    out['spread_bps'] = out['vpin_regime'].map({k: v['spread_bps'] for k, v in regime_params.items()})
    out['size_frac']  = out['vpin_regime'].map({k: v['size_frac'] for k, v in regime_params.items()})
    return out


validation_summary = validate_vpin_predictiveness(vpin_df)
vpin_with_response = compute_mm_toxicity_response(vpin_df)

def plot_vpin_dashboard(
    vpin_df: pd.DataFrame,
    trades: pd.DataFrame
) -> None:
    """
    Three-panel VPIN dashboard: price, VPIN timeline, and regime distribution.

    Parameters
    ----------
    vpin_df : pd.DataFrame
        VPIN output with regime classifications.
    trades : pd.DataFrame
        Raw trade data for price history.
    """
    fig, axes = plt.subplots(3, 1, figsize=(14, 12), sharex=False)

    # Panel 1: Price history with toxic episodes shaded
    axes[0].plot(trades['trade_id'], trades['mid_price'], color='steelblue', linewidth=0.5)
    toxic_mask = trades['is_toxic_episode'] == 1
    axes[0].fill_between(trades['trade_id'], trades['mid_price'].min(), trades['mid_price'].max(),
                          where=toxic_mask, alpha=0.2, color='red', label='Toxic Episode')
    axes[0].set_ylabel('Mid Price')
    axes[0].set_title('Price History — Red = Known Toxic Episodes')
    axes[0].legend()

    # Panel 2: VPIN timeline
    valid = vpin_df.dropna(subset=['vpin'])
    axes[1].plot(valid['bucket_id'], valid['vpin'], color='black', linewidth=1.5)
    axes[1].fill_between(valid['bucket_id'], valid['vpin'], VPIN_HIGH_THRESH,
                          where=(valid['vpin'] >= VPIN_HIGH_THRESH), alpha=0.3, color='red', label=f'High toxicity (>{VPIN_HIGH_THRESH})')
    axes[1].fill_between(valid['bucket_id'], valid['vpin'], VPIN_LOW_THRESH,
                          where=(valid['vpin'] <= VPIN_LOW_THRESH), alpha=0.3, color='green', label=f'Low toxicity (<{VPIN_LOW_THRESH})')
    axes[1].axhline(VPIN_HIGH_THRESH, color='red',   linewidth=0.8, linestyle='--')
    axes[1].axhline(VPIN_LOW_THRESH,  color='green', linewidth=0.8, linestyle='--')
    axes[1].set_ylabel('VPIN')
    axes[1].set_title('VPIN (Volume-Synchronized Probability of Informed Trading)')
    axes[1].set_xlabel('Bucket ID')
    axes[1].legend()

    # Panel 3: Spread response
    if 'spread_bps' in vpin_df.columns:
        axes[2].bar(vpin_df['bucket_id'], vpin_df['spread_bps'],
                    color=vpin_df['vpin_regime'].map({'low': 'green', 'medium': 'gold', 'high': 'red', None: 'grey'}),
                    alpha=0.7, width=1.0)
        axes[2].set_ylabel('Quoted Spread (bps)')
        axes[2].set_title('MM Quoted Spread by Toxicity Regime')
        axes[2].set_xlabel('Bucket ID')
    else:
        axes[2].set_visible(False)

    plt.tight_layout()
    plt.show()


plot_vpin_dashboard(vpin_with_response, trades)

def export_toxicity_data(
    vpin_df: pd.DataFrame,
    validation_summary: pd.DataFrame
) -> None:
    """
    Export VPIN and validation data.

    Parameters
    ----------
    vpin_df : pd.DataFrame
        VPIN per bucket with regime and MM response parameters.
    validation_summary : pd.DataFrame
        VPIN regime vs future price move analysis.
    """
    vpin_df.to_csv('vpin_toxicity.csv', index=False)
    validation_summary.to_csv('vpin_validation.csv')
    print('Exported: vpin_toxicity.csv')
    print('Exported: vpin_validation.csv')


export_toxicity_data(vpin_with_response, validation_summary)