pip install pandas numpy scipy matplotlib seaborn

import logging
import random
from collections import deque
import time
import math
from typing import Dict, Any, Callable

import pandas as pd
import numpy as np
from scipy.stats import norm
import matplotlib.pyplot as plt
import seaborn as sns

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def create_market_making_state(initial_cash: float, initial_base_asset: float, current_price: float,
                               spread_bps: float, max_inventory_delta: float, pnl_history_size: int = 100) -> Dict[str, Any]:
    """
    Initializes the state dictionary for a crypto market making strategy.

    Parameters
    ----------
    initial_cash : float
        The starting cash balance (quote asset).
    initial_base_asset : float
        The starting quantity of the base asset.
    current_price : float
        The current market price of the base asset.
    spread_bps : float
        The base spread in basis points (e.g., 5 for 0.05%).
    max_inventory_delta : float
        The maximum absolute inventory delta allowed.
    pnl_history_size : int, optional
        The number of past PnL values to store, defaults to 100. (Note: For full history plotting, this is effectively ignored as `total_pnl_history` is a list.)

    Returns
    -------
    Dict[str, Any]
        An initialized state dictionary containing market making parameters.

    Examples
    --------
    >>> state = create_market_making_state(initial_cash=10000, initial_base_asset=0.5, current_price=20000, spread_bps=5, max_inventory_delta=0.1)
    >>> print(state['current_cash'])
    10000.0
    """
    state = {
        'current_cash': initial_cash,
        'current_base_asset': initial_base_asset,
        'inventory_delta': initial_base_asset, # Net quantity of base asset
        'current_price': current_price,
        'spread_bps': spread_bps,
        'max_inventory_delta': max_inventory_delta,
        'trade_count': 0,
        'long_trades': 0,
        'short_trades': 0,
        'realized_pnl': 0.0,
        'unrealized_pnl': 0.0,
        'total_pnl_history': [0.0], # Changed from deque to list to track full history
        'inventory_delta_history': [initial_base_asset],
        'price_history': [current_price],
        'balance_history': [], # total value in quote asset
        'time_step': 0
    }
    # Calculate initial total balance
    state['total_balance'] = initial_cash + (initial_base_asset * current_price)
    state['balance_history'].append(state['total_balance'])

    logger.info(f"Market making state initialized with cash={initial_cash}, base_asset={initial_base_asset}, price={current_price}")
    return state

def place_orders(state: Dict[str, Any], order_size: float) -> Dict[str, Any]:
    """
    Simulates placing bid and ask orders and potential fills.

    Parameters
    ----------
    state : Dict[str, Any]
        The current market making state dictionary.
    order_size : float
        The quantity of the base asset for each bid/ask order.

    Returns
    -------
    Dict[str, Any]
        The updated state dictionary after simulating order placement and fills.

    Examples
    --------
    >>> state = create_market_making_state(initial_cash=10000, initial_base_asset=0.5, current_price=20000, spread_bps=5, max_inventory_delta=0.1)
    >>> updated_state = place_orders(state, order_size=0.01)
    >>> isinstance(updated_state, dict)
    True
    """
    current_price = state['current_price']
    spread = current_price * (state['spread_bps'] / 10000)

    bid_price = current_price - spread / 2
    ask_price = current_price + spread / 2

    logger.debug(f"Placing orders: Bid at {bid_price:.2f}, Ask at {ask_price:.2f}")

    # Simulate fills (simplified: assume 50/50 chance of buy/sell fill if within spread)
    # In a real scenario, this would depend on market conditions and order book depth.
    trade_occurred = False
    if random.random() < 0.4: # Simulate a bid fill (someone bought from our ask)
        fill_qty = order_size * random.uniform(0.5, 1.0) # Partial or full fill
        state['current_cash'] += fill_qty * ask_price
        state['current_base_asset'] -= fill_qty
        state['inventory_delta'] -= fill_qty
        state['realized_pnl'] += fill_qty * (ask_price - bid_price) # Simplified PnL from spread capture
        state['trade_count'] += 1
        state['short_trades'] += 1
        logger.info(f"Sold {fill_qty:.4f} base asset at {ask_price:.2f}. New inventory delta: {state['inventory_delta']:.4f}")
        trade_occurred = True

    if random.random() < 0.4: # Simulate an ask fill (someone sold to our bid)
        fill_qty = order_size * random.uniform(0.5, 1.0)
        state['current_cash'] -= fill_qty * bid_price
        state['current_base_asset'] += fill_qty
        state['inventory_delta'] += fill_qty
        state['realized_pnl'] += fill_qty * (ask_price - bid_price) # Simplified PnL from spread capture
        state['trade_count'] += 1
        state['long_trades'] += 1
        logger.info(f"Bought {fill_qty:.4f} base asset at {bid_price:.2f}. New inventory delta: {state['inventory_delta']:.4f}")
        trade_occurred = True

    if not trade_occurred:
        logger.debug("No fills occurred in this step.")

    return state

def adjust_inventory_risk(state: Dict[str, Any]) -> Dict[str, Any]:
    """
    Adjusts inventory to manage risk by rebalancing if inventory_delta exceeds limits.

    Parameters
    ----------
    state : Dict[str, Any]
        The current market making state dictionary.

    Returns
    -------
    Dict[str, Any]
        The updated state dictionary after any inventory adjustments.

    Examples
    --------
    >>> state = create_market_making_state(initial_cash=10000, initial_base_asset=0.5, current_price=20000, spread_bps=5, max_inventory_delta=0.1)
    >>> state['inventory_delta'] = 0.15 # Artificially increase delta to trigger rebalance
    >>> updated_state = adjust_inventory_risk(state)
    >>> isinstance(updated_state, dict)
    True
    """
    current_delta = state['inventory_delta']
    max_delta = state['max_inventory_delta']
    current_price = state['current_price']

    if abs(current_delta) > max_delta:
        adjustment_amount = abs(current_delta) - max_delta
        if current_delta > 0: # Inventory is too long, need to sell
            state['current_cash'] += adjustment_amount * current_price
            state['current_base_asset'] -= adjustment_amount
            state['inventory_delta'] -= adjustment_amount
            state['realized_pnl'] -= adjustment_amount * (current_price * (state['spread_bps'] / 10000)) # Cost of rebalancing
            logger.warning(f"Rebalancing: Sold {adjustment_amount:.4f} base asset at {current_price:.2f} due to long inventory. New delta: {state['inventory_delta']:.4f}")
        else: # Inventory is too short, need to buy
            state['current_cash'] -= adjustment_amount * current_price
            state['current_base_asset'] += adjustment_amount
            state['inventory_delta'] += adjustment_amount
            state['realized_pnl'] -= adjustment_amount * (current_price * (state['spread_bps'] / 10000)) # Cost of rebalancing
            logger.warning(f"Rebalancing: Bought {adjustment_amount:.4f} base asset at {current_price:.2f} due to short inventory. New delta: {state['inventory_delta']:.4f}")
    else:
        logger.debug(f"Inventory delta {current_delta:.4f} is within limits [{ -max_delta:.4f}, {max_delta:.4f}]")

    return state

def track_metrics(state: Dict[str, Any]) -> Dict[str, Any]:
    """
    Updates various metrics within the market making state dictionary.

    Parameters
    ----------
    state : Dict[str, Any]
        The current market making state dictionary.

    Returns
    -------
    Dict[str, Any]
        The updated state dictionary with new metric values and histories.

    Examples
    --------
    >>> state = create_market_making_state(initial_cash=10000, initial_base_asset=0.5, current_price=20000, spread_bps=5, max_inventory_delta=0.1)
    >>> updated_state = track_metrics(state)
    >>> isinstance(updated_state, dict)
    True
    """
    current_price = state['current_price']
    current_cash = state['current_cash']
    current_base_asset = state['current_base_asset']

    # Calculate unrealized PnL (from current base asset holdings)
    # Assuming initial_base_asset and initial_cash were used to define initial state for PnL calculation basis
    initial_total_value = state['total_balance'] # Total value at the start
    current_total_value = current_cash + (current_base_asset * current_price)
    state['unrealized_pnl'] = current_total_value - initial_total_value - state['realized_pnl'] # Total PnL minus realized

    # Update total balance history
    state['balance_history'].append(current_total_value)

    # Update PnL history
    state['total_pnl_history'].append(state['realized_pnl'] + state['unrealized_pnl'])

    # Update inventory delta history
    state['inventory_delta_history'].append(state['inventory_delta'])

    # Update price history
    state['price_history'].append(current_price)

    state['time_step'] += 1

    logger.debug(f"Metrics tracked for time step {state['time_step']}. Total PnL: {state['total_pnl_history'][-1]:.2f}")
    return state

def exponential_backoff_retry(func: Callable, max_retries: int = 5, initial_delay: float = 1.0, *args, **kwargs) -> Any:
    """
    Retries a function with exponential backoff.

    Parameters
    ----------
    func : Callable
        The function to be executed and retried.
    max_retries : int, optional
        The maximum number of times to retry the function, defaults to 5.
    initial_delay : float, optional
        The initial delay in seconds before the first retry, defaults to 1.0.
    *args
        Positional arguments to pass to the `func`.
    **kwargs
        Keyword arguments to pass to the `func`.

    Returns
    -------
    Any
        The result of the successful `func` execution.

    Raises
    -------
    Exception
        If the function fails after all retries, the last exception is re-raised.

    Examples
    --------
    >>> call_count = 0
    >>> def flaky_function():
    ...     nonlocal call_count
    ...     call_count += 1
    ...     if call_count < 3:
    ...         raise ValueError("Temporary error")
    ...     return "Success"
    >>> # result = exponential_backoff_retry(flaky_function, max_retries=3, initial_delay=0.1)
    >>> # print(result)
    """
    delay = initial_delay
    for i in range(max_retries):
        try:
            logger.debug(f"Attempt {i + 1}/{max_retries} for function {func.__name__}")
            return func(*args, **kwargs)
        except Exception as e:
            logger.warning(f"Attempt {i + 1} failed for {func.__name__}: {e}. Retrying in {delay:.2f} seconds...")
            time.sleep(delay)
            delay *= 2  # Exponential backoff
            jitter = random.uniform(0.8 * initial_delay, 1.2 * initial_delay) # Add jitter
            delay += jitter
    logger.error(f"Function {func.__name__} failed after {max_retries} retries.")
    raise Exception(f"Failed after {max_retries} retries")

def simulate_market_making(initial_state: Dict[str, Any],
                           num_steps: int,
                           order_size: float,
                           volatility_annual: float,
                           time_step_days: float = 1/365,
                           drift: float = 0) -> Dict[str, Any]:
    """
    Simulates a market making strategy over a specified number of time steps.

    Parameters
    ----------
    initial_state : Dict[str, Any]
        The initial market making state dictionary.
    num_steps : int
        The total number of simulation steps.
    order_size : float
        The quantity of the base asset for each bid/ask order.
    volatility_annual : float
        The annualized volatility for price simulation.
    time_step_days : float, optional
        The time duration of each step in days, defaults to 1/365.
    drift : float, optional
        The average rate of return (drift) of the asset, defaults to 0.

    Returns
    -------
    Dict[str, Any]
        The final state dictionary after the simulation, containing all historical data.

    Examples
    --------
    >>> initial_state = create_market_making_state(initial_cash=10000, initial_base_asset=0.5, current_price=20000, spread_bps=5, max_inventory_delta=0.1)
    >>> final_state = simulate_market_making(initial_state, num_steps=10, order_size=0.01, volatility_annual=0.7)
    >>> isinstance(final_state, dict)
    True
    """
    state = initial_state.copy()
    logger.info(f"Starting market making simulation for {num_steps} steps.")

    for step in range(num_steps):
        logger.debug(f"Simulation Step {step + 1}/{num_steps}")

        # 1. Simulate price movement
        new_price = exponential_backoff_retry(simulate_price_movement, 5, 1.0, state, volatility_annual, time_step_days, drift)
        state['current_price'] = new_price

        # 2. Place orders and simulate fills
        state = exponential_backoff_retry(place_orders, 5, 1.0, state, order_size)

        # 3. Adjust inventory risk if necessary
        state = exponential_backoff_retry(adjust_inventory_risk, 5, 1.0, state)

        # 4. Track metrics
        state = exponential_backoff_retry(track_metrics, 5, 1.0, state)

    logger.info("Market making simulation completed.")
    return state

# Define simulation parameters
INITIAL_CASH = 100000.0 # Initial capital in quote asset (e.g., USDT)
INITIAL_BASE_ASSET = 1.0 # Initial quantity of base asset (e.g., BTC)
INITIAL_PRICE = 30000.0 # Starting price of the base asset
SPREAD_BPS = 10 # Bid-ask spread in basis points (10 bps = 0.1%)
MAX_INVENTORY_DELTA = 0.5 # Maximum allowed absolute inventory delta (0.5 BTC)
SIMULATION_STEPS = 500 # Number of simulation steps (e.g., hours, minutes)
ORDER_SIZE = 0.05 # Quantity of base asset per order
VOLATILITY_ANNUAL = 0.8 # Annualized volatility (e.g., 80% for crypto)
TIME_STEP_DAYS = 1 / (365 * 24) # Each step represents one hour
DRIFT = 0.05 # Annualized drift (e.g., 5% expected annual return)

logger.info("Initializing simulation state...")
initial_mm_state = create_market_making_state(
    initial_cash=INITIAL_CASH,
    initial_base_asset=INITIAL_BASE_ASSET,
    current_price=INITIAL_PRICE,
    spread_bps=SPREAD_BPS,
    max_inventory_delta=MAX_INVENTORY_DELTA
)

logger.info("Starting market making simulation...")
final_mm_state = simulate_market_making(
    initial_state=initial_mm_state,
    num_steps=SIMULATION_STEPS,
    order_size=ORDER_SIZE,
    volatility_annual=VOLATILITY_ANNUAL,
    time_step_days=TIME_STEP_DAYS,
    drift=DRIFT
)

print("\n--- Simulation Summary ---")
print(f"Final Cash: {final_mm_state['current_cash']:.2f}")
print(f"Final Base Asset: {final_mm_state['current_base_asset']:.4f}")
print(f"Final Price: {final_mm_state['current_price']:.2f}")
print(f"Total Trades: {final_mm_state['trade_count']}")
print(f"Realized PnL: {final_mm_state['realized_pnl']:.2f}")
print(f"Unrealized PnL: {final_mm_state['unrealized_pnl']:.2f}")
print(f"Total PnL: {(final_mm_state['realized_pnl'] + final_mm_state['unrealized_pnl']):.2f}")

import matplotlib.pyplot as plt
import seaborn as sns

# Convert histories to DataFrame for easier plotting
df_history = pd.DataFrame({
    'step': range(len(final_mm_state['price_history'])),
    'price': final_mm_state['price_history'],
    'total_pnl': final_mm_state['total_pnl_history'], # Now a list
    'inventory_delta': final_mm_state['inventory_delta_history']
})

plt.figure(figsize=(15, 12))

# Plot 1: Price Movement
plt.subplot(3, 1, 1) # 3 rows, 1 column, 1st plot
sns.lineplot(x='step', y='price', data=df_history)
plt.title('Simulated Asset Price Movement')
plt.xlabel('Simulation Step')
plt.ylabel('Price')
plt.grid(True)

# Plot 2: Total PnL Over Time
plt.subplot(3, 1, 2) # 3 rows, 1 column, 2nd plot
sns.lineplot(x='step', y='total_pnl', data=df_history, color='green')
plt.title('Cumulative Total PnL Over Time')
plt.xlabel('Simulation Step')
plt.ylabel('Total PnL (Quote Asset)')
plt.grid(True)

# Plot 3: Inventory Delta Over Time
plt.subplot(3, 1, 3) # 3 rows, 1 column, 3rd plot
sns.lineplot(x='step', y='inventory_delta', data=df_history, color='red')
plt.axhspan(-MAX_INVENTORY_DELTA, MAX_INVENTORY_DELTA, color='grey', alpha=0.2, label='Max Inventory Delta Limit')
plt.title('Inventory Delta Over Time')
plt.xlabel('Simulation Step')
plt.ylabel('Inventory Delta (Base Asset)')
plt.legend()
plt.grid(True)

plt.tight_layout()
plt.show()