!pip install numpy pandas matplotlib seaborn --quiet

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from dataclasses import dataclass, field
from typing import Tuple, Optional
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('muted')
print('Imports ready.')

# ── MARKET PARAMETERS ──────────────────────────────────────────────────────────
BASE_QUOTE_SIZE    = 0.10    # BTC — default size when inventory is neutral
MAX_INVENTORY      = 1.00    # BTC — hard inventory limit (each side)
INVENTORY_HALF_CAP = 0.50    # BTC — size starts shrinking at this inventory level
MIN_QUOTE_SIZE     = 0.01    # BTC — minimum quote size regardless of adjustments

VOLATILITY_SCALAR  = 0.30    # how strongly vol affects size: 0=no effect, 1=full effect
SPREAD_SCALAR      = 0.20    # how strongly spread width affects size

SIMULATION_STEPS   = 5_000   # number of simulated ticks
# ────────────────────────────────────────────────────────────────────────────────

@dataclass
class MarketState:
    """
    Snapshot of the current market and MM position used for sizing decisions.

    Attributes
    ----------
    mid_price : float
        Current market mid-price.
    realized_vol : float
        Recent realized volatility as a decimal fraction (e.g., 0.02 = 2%).
    current_spread : float
        Current best bid-ask spread as a decimal fraction.
    inventory : float
        Current net inventory in base asset (positive = long, negative = short).
    """
    mid_price:      float
    realized_vol:   float
    current_spread: float
    inventory:      float


def compute_inventory_factor(
    inventory: float,
    max_inventory: float,
    half_cap: float
) -> Tuple[float, float]:
    """
    Compute bid and ask sizing factors based on current inventory level.

    When inventory is positive (long), we reduce ask size and bid size:
    - Reduce bid size (don't want more longs)
    - Keep or slightly increase ask size (want to sell to reduce inventory)

    Parameters
    ----------
    inventory : float
        Current inventory in base asset.
    max_inventory : float
        Hard limit — beyond this, bid_factor → 0 (stop buying) or ask_factor → 0 (stop selling).
    half_cap : float
        Inventory level where size starts to taper. Below this: factor=1.0.

    Returns
    -------
    Tuple[float, float]
        (bid_factor, ask_factor) in [0, 1].

    Notes
    -----
    The inventory factor is computed symmetrically: positive inventory penalizes
    the bid side (reduces willingness to buy more), negative penalizes the ask side.
    At max_inventory the penalized side drops to 0 (quoting stopped on that side).
    """
    inv_norm = inventory / max_inventory  # normalized to [-1, 1]

    # Bid factor: decays as inventory grows positive
    if inventory >= 0:
        half_norm  = half_cap / max_inventory
        bid_factor = max(0.0, 1.0 - (inv_norm / 1.0) * (1.0 / (1.0 - half_norm + 1e-8)))
        ask_factor = min(1.0, 1.0 + inv_norm * 0.2)  # slightly more aggressive selling
    else:
        half_norm  = half_cap / max_inventory
        ask_factor = max(0.0, 1.0 - (-inv_norm / 1.0) * (1.0 / (1.0 - half_norm + 1e-8)))
        bid_factor = min(1.0, 1.0 + (-inv_norm) * 0.2)

    return round(max(0.0, bid_factor), 4), round(max(0.0, ask_factor), 4)


def compute_volatility_factor(
    realized_vol: float,
    reference_vol: float = 0.02,
    scalar: float = 0.30
) -> float:
    """
    Compute size scaling factor based on current vs reference volatility.

    Parameters
    ----------
    realized_vol : float
        Recent realized volatility (decimal fraction, e.g., 0.02 = 2%).
    reference_vol : float
        Baseline volatility at which size factor = 1.0.
    scalar : float
        How aggressively to shrink size when vol exceeds reference.
        0 = no adjustment, 1 = full adjustment.

    Returns
    -------
    float
        Volatility size factor in (0, 1].

    Notes
    -----
    High volatility increases adverse selection risk: a large, slow quote is more
    likely to be picked off by informed traders. Shrinking size in high-vol periods
    is therefore both a risk management and alpha protection measure.
    """
    if realized_vol <= 0 or reference_vol <= 0:
        return 1.0
    vol_ratio   = reference_vol / realized_vol  # <1 when vol is high
    raw_factor  = vol_ratio ** scalar
    return round(min(1.0, max(0.1, raw_factor)), 4)


def compute_spread_factor(
    current_spread: float,
    reference_spread: float = 0.0005,
    scalar: float = 0.20
) -> float:
    """
    Compute size scaling factor based on spread width.

    Wider spreads provide more compensation per fill — can afford slightly larger quotes.

    Parameters
    ----------
    current_spread : float
        Current bid-ask spread as a decimal fraction.
    reference_spread : float
        Baseline spread at which factor = 1.0.
    scalar : float
        Sensitivity to spread changes.

    Returns
    -------
    float
        Spread size factor, bounded to [0.5, 1.5].
    """
    if current_spread <= 0 or reference_spread <= 0:
        return 1.0
    spread_ratio = current_spread / reference_spread
    factor = 1.0 + scalar * (spread_ratio - 1.0)
    return round(min(1.5, max(0.5, factor)), 4)


def compute_quote_sizes(
    state: MarketState,
    base_size: float = BASE_QUOTE_SIZE,
    max_inventory: float = MAX_INVENTORY,
    half_cap: float = INVENTORY_HALF_CAP,
    min_size: float = MIN_QUOTE_SIZE,
    vol_scalar: float = VOLATILITY_SCALAR,
    spread_scalar: float = SPREAD_SCALAR
) -> dict:
    """
    Compute final bid and ask quote sizes given current market state.

    Combines inventory, volatility, and spread adjustments.

    Parameters
    ----------
    state : MarketState
        Current market and position state.
    base_size : float
        Base quote size in base asset units.
    max_inventory : float
        Hard inventory limit.
    half_cap : float
        Inventory level where tapering begins.
    min_size : float
        Minimum allowed quote size.
    vol_scalar : float
        Volatility adjustment strength.
    spread_scalar : float
        Spread-width adjustment strength.

    Returns
    -------
    dict
        Keys: bid_size, ask_size, bid_factor, ask_factor, vol_factor, spread_factor.
    """
    bid_inv_f, ask_inv_f = compute_inventory_factor(state.inventory, max_inventory, half_cap)
    vol_f    = compute_volatility_factor(state.realized_vol, scalar=vol_scalar)
    spread_f = compute_spread_factor(state.current_spread, scalar=spread_scalar)

    combined = vol_f * spread_f
    raw_bid  = base_size * bid_inv_f * combined
    raw_ask  = base_size * ask_inv_f * combined

    # Round to exchange minimum tick (0.001 BTC assumed)
    bid_size = max(min_size, round(raw_bid / 0.001) * 0.001)
    ask_size = max(min_size, round(raw_ask / 0.001) * 0.001)

    # If inventory has hit the hard limit, stop quoting on that side
    if state.inventory >= max_inventory:
        bid_size = 0.0
    if state.inventory <= -max_inventory:
        ask_size = 0.0

    return {
        'bid_size': bid_size, 'ask_size': ask_size,
        'bid_inv_factor': bid_inv_f, 'ask_inv_factor': ask_inv_f,
        'vol_factor': vol_f, 'spread_factor': spread_f
    }

def generate_synthetic_market(
    n_steps: int,
    start_price: float = 50_000,
    base_vol: float = 0.02,
    base_spread: float = 0.0005
) -> pd.DataFrame:
    """
    Generate synthetic market tick data with realistic vol and spread dynamics.

    Parameters
    ----------
    n_steps : int
        Number of simulation ticks.
    start_price : float
        Initial mid-price.
    base_vol : float
        Baseline daily volatility (decimal fraction).
    base_spread : float
        Baseline bid-ask spread (decimal fraction).

    Returns
    -------
    pd.DataFrame
        Columns: mid_price, realized_vol, spread, buy_prob.
        buy_prob: probability that next order is a buy (drives inventory dynamics).
    """
    np.random.seed(42)
    ticks_per_day  = 288  # 5-minute ticks
    tick_vol       = base_vol / np.sqrt(ticks_per_day)

    mid_prices = [start_price]
    for _ in range(n_steps - 1):
        ret = np.random.normal(0, tick_vol)
        mid_prices.append(mid_prices[-1] * (1 + ret))

    mid_prices = np.array(mid_prices)

    # Realized vol: 20-tick rolling std of returns
    log_rets = np.log(mid_prices[1:] / mid_prices[:-1])
    realized_vol = pd.Series(log_rets).rolling(20).std().fillna(tick_vol).values
    realized_vol = np.append(tick_vol, realized_vol) * np.sqrt(ticks_per_day)  # annualize to daily

    # Spread: widens when vol is high
    spread = base_spread * (1 + 2 * realized_vol / base_vol)

    # Buy probability: slight autocorrelation (momentum)
    buy_prob = 0.5 + 0.1 * np.sign(np.append(0, log_rets))

    return pd.DataFrame({
        'mid_price':    mid_prices,
        'realized_vol': realized_vol.clip(0.005, 0.20),
        'spread':       spread.clip(0.0001, 0.005),
        'buy_prob':     buy_prob,
    })


market_data = generate_synthetic_market(SIMULATION_STEPS)
print(f'Synthetic market: {len(market_data)} ticks')
print(market_data.describe().round(5))

def simulate_quote_sizing(
    market_data: pd.DataFrame,
    base_size: float = BASE_QUOTE_SIZE,
    max_inventory: float = MAX_INVENTORY,
    half_cap: float = INVENTORY_HALF_CAP,
    fill_prob_per_level: float = 0.15
) -> pd.DataFrame:
    """
    Simulate a market maker running the quote sizing logic over tick data.

    At each tick: compute quote sizes, randomly simulate fills based on
    fill probability and buy_prob, update inventory.

    Parameters
    ----------
    market_data : pd.DataFrame
        Output of generate_synthetic_market().
    base_size : float
        Base quote size in BTC.
    max_inventory : float
        Inventory hard limit.
    half_cap : float
        Inventory taper start.
    fill_prob_per_level : float
        Probability per tick that each side gets a fill (simplified).

    Returns
    -------
    pd.DataFrame
        Simulation history with inventory, quote sizes, and PnL.
    """
    records = []
    inventory = 0.0
    cash      = 0.0

    for i, row in market_data.iterrows():
        state = MarketState(
            mid_price    = row['mid_price'],
            realized_vol = row['realized_vol'],
            current_spread = row['spread'],
            inventory    = inventory
        )
        sizes = compute_quote_sizes(state, base_size, max_inventory, half_cap)

        # Quote prices (half-spread each side)
        half_spread = row['mid_price'] * row['spread'] / 2
        bid_price   = row['mid_price'] - half_spread
        ask_price   = row['mid_price'] + half_spread

        # Simulate fills
        bid_filled = sizes['bid_size'] > 0 and np.random.rand() < fill_prob_per_level * row['buy_prob'] * 2
        ask_filled = sizes['ask_size'] > 0 and np.random.rand() < fill_prob_per_level * (1 - row['buy_prob']) * 2

        if bid_filled:
            inventory += sizes['bid_size']
            cash      -= sizes['bid_size'] * bid_price
        if ask_filled:
            inventory -= sizes['ask_size']
            cash      += sizes['ask_size'] * ask_price

        # Mark-to-market PnL = cash + inventory × mid_price
        mtm_pnl = cash + inventory * row['mid_price']

        records.append({
            'tick':         i,
            'mid_price':    row['mid_price'],
            'inventory':    inventory,
            'cash':         cash,
            'mtm_pnl':      mtm_pnl,
            'bid_size':     sizes['bid_size'],
            'ask_size':     sizes['ask_size'],
            'vol_factor':   sizes['vol_factor'],
            'bid_inv_factor': sizes['bid_inv_factor'],
        })

    return pd.DataFrame(records)


sim = simulate_quote_sizing(market_data)
print(f'Simulation complete.')
print(f'Final inventory: {sim["inventory"].iloc[-1]:.4f} BTC')
print(f'Final MtM PnL:  ${sim["mtm_pnl"].iloc[-1]:.2f}')
print(f'Max inventory:   {sim["inventory"].abs().max():.4f} BTC')

def plot_quote_sizing_simulation(sim: pd.DataFrame) -> None:
    """
    Four-panel dashboard showing inventory, quote sizes, factors, and PnL.

    Parameters
    ----------
    sim : pd.DataFrame
        Output of simulate_quote_sizing().
    """
    fig, axes = plt.subplots(4, 1, figsize=(14, 16), sharex=True)

    axes[0].plot(sim['tick'], sim['mid_price'], color='steelblue', linewidth=0.8)
    axes[0].set_ylabel('Mid Price')
    axes[0].set_title('Mid Price')

    axes[1].fill_between(sim['tick'], sim['inventory'], alpha=0.4,
                          color=np.where(sim['inventory'] >= 0, 'green', 'red').tolist()[0])
    axes[1].plot(sim['tick'], sim['inventory'], linewidth=0.6, color='black')
    axes[1].axhline(0, color='black', linewidth=0.8, linestyle='--')
    axes[1].axhline( MAX_INVENTORY,  color='red', linewidth=0.8, linestyle=':')
    axes[1].axhline(-MAX_INVENTORY,  color='red', linewidth=0.8, linestyle=':')
    axes[1].set_ylabel('Inventory (BTC)')
    axes[1].set_title('Inventory — Red Lines = Hard Limits')

    axes[2].plot(sim['tick'], sim['bid_size'], label='Bid Size', color='green', linewidth=0.6, alpha=0.8)
    axes[2].plot(sim['tick'], sim['ask_size'], label='Ask Size', color='red',   linewidth=0.6, alpha=0.8)
    axes[2].plot(sim['tick'], sim['vol_factor'] * BASE_QUOTE_SIZE, label='Vol-scaled base',
                  color='black', linewidth=0.6, linestyle='--')
    axes[2].set_ylabel('Quote Size (BTC)')
    axes[2].set_title('Bid and Ask Quote Sizes')
    axes[2].legend()

    axes[3].plot(sim['tick'], sim['mtm_pnl'], color='gold', linewidth=1.0)
    axes[3].axhline(0, color='black', linewidth=0.8, linestyle='--')
    axes[3].fill_between(sim['tick'], sim['mtm_pnl'], 0,
                          where=(sim['mtm_pnl'] >= 0), alpha=0.2, color='green')
    axes[3].fill_between(sim['tick'], sim['mtm_pnl'], 0,
                          where=(sim['mtm_pnl'] < 0), alpha=0.2, color='red')
    axes[3].set_ylabel('Mark-to-Market PnL ($)')
    axes[3].set_title('Mark-to-Market PnL')
    axes[3].set_xlabel('Tick')

    plt.tight_layout()
    plt.show()


plot_quote_sizing_simulation(sim)

def export_quote_sizing(sim: pd.DataFrame) -> None:
    """
    Export simulation results to CSV.

    Parameters
    ----------
    sim : pd.DataFrame
        Quote sizing simulation output.
    """
    sim.to_csv('quote_sizing_simulation.csv', index=False)
    print(f'Exported quote_sizing_simulation.csv ({len(sim)} ticks)')


export_quote_sizing(sim)