pip install numpy pandas scipy matplotlib seaborn loguru tqdm

# Standard library imports
import collections
import math
import random
import time
import logging

# Third-party imports
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.stats import norm
from loguru import logger
from tqdm.notebook import tqdm

# Configure loguru logger
logger.remove()
logger.add(lambda msg: logging.info(msg.strip()), level="INFO")
logger.add(lambda msg: logging.debug(msg.strip()), level="DEBUG")
logger.add(lambda msg: logging.warning(msg.strip()), level="WARNING")

def create_market_state(
    initial_cash: float,
    initial_inventory: int,
    initial_mid_price: float,
    tick_size: float = 0.01,
    max_inventory: int = 100,
    min_inventory: int = -100,
    risk_aversion: float = 0.001,
    volatility: float = 0.01,
    kappa: float = 0.01,
    lambda_val: float = 1.0,
    order_book_impact: float = 0.05
) -> dict:
    """
    Initializes the market maker's state dictionary.

    Parameters
    ----------
    initial_cash : float
        The starting cash balance for the market maker.
    initial_inventory : int
        The starting inventory of the asset.
    initial_mid_price : float
        The initial mid-price of the asset.
    tick_size : float, optional
        The minimum price increment, defaults to 0.01.
    max_inventory : int, optional
        The maximum allowed inventory for the asset, defaults to 100.
    min_inventory : int, optional
        The minimum allowed inventory, defaults to -100.
    risk_aversion : float, optional
        A parameter reflecting the market maker's aversion to inventory risk, defaults to 0.001.
    volatility : float, optional
        The assumed market volatility for spread calculation, defaults to 0.01.
    kappa : float, optional
        A parameter related to order book resilience/arrival rate, defaults to 0.01.
    lambda_val : float, optional
        A parameter related to order book depth/intensity, defaults to 1.0.
    order_book_impact : float, optional
        How much order book depth or imbalance influences the spread, defaults to 0.05.

    Returns
    -------
    dict
        An initialized state dictionary containing all market maker parameters.
    """
    logger.info("Initializing market state.")
    state = {
        'cash': initial_cash,
        'inventory': initial_inventory,
        'mid_price': initial_mid_price,
        'tick_size': tick_size,
        'max_inventory': max_inventory,
        'min_inventory': min_inventory,
        'risk_aversion': risk_aversion,
        'volatility': volatility,
        'kappa': kappa,
        'lambda_val': lambda_val,
        'order_book_impact': order_book_impact,
        'time': 0
    }
    logger.debug(f"Market state initialized: {state}")
    return state

def generate_market_data(state: dict, num_steps: int) -> pd.DataFrame:
    """
    Simulates market data including mid-price, bid/ask prices, and order book depth.

    Parameters
    ----------
    state : dict
        The current market state dictionary.
    num_steps : int
        The number of simulation steps to generate data for.

    Returns
    -------
    pd.DataFrame
        A DataFrame containing simulated market data (time, mid_price, bid_price, ask_price, bid_volume, ask_volume)
        for `num_steps`.

    Examples
    --------
    >>> initial_state = create_market_state(10000, 0, 100.0)
    >>> market_data = generate_market_data(initial_state, 10)
    >>> print(market_data.head())
    """
    logger.info(f"Generating {num_steps} steps of market data.")
    mid_prices = [state['mid_price']]
    bid_prices = []
    ask_prices = []
    bid_volumes = []
    ask_volumes = []
    times = [0]

    current_mid_price = state['mid_price']
    tick_size = state['tick_size']

    for i in range(1, num_steps):
        # Simulate mid-price as a random walk
        current_mid_price += np.random.normal(0, state['volatility']) * current_mid_price
        current_mid_price = round(current_mid_price / tick_size) * tick_size
        mid_prices.append(current_mid_price)
        times.append(i)

        # Simulate some base bid/ask spread and depth
        base_spread = 2 * tick_size # A small base spread
        bid = current_mid_price - base_spread / 2 - random.uniform(0, tick_size)
        ask = current_mid_price + base_spread / 2 + random.uniform(0, tick_size)

        bid = round(bid / tick_size) * tick_size
        ask = round(ask / tick_size) * tick_size

        bid_prices.append(bid)
        ask_prices.append(ask)

        # Simulate order book depth with some randomness
        bid_volumes.append(max(1, int(np.random.normal(50, 10))))
        ask_volumes.append(max(1, int(np.random.normal(50, 10))))

    df = pd.DataFrame({
        'time': times,
        'mid_price': mid_prices,
        'bid_price': bid_prices + [0] * (num_steps - len(bid_prices)), # Adjust for initial 0th step, fill with 0, will be overwritten if used with real data
        'ask_price': ask_prices + [0] * (num_steps - len(ask_prices)), # Adjust similarly
        'bid_volume': bid_volumes + [0] * (num_steps - len(bid_volumes)),
        'ask_volume': ask_volumes + [0] * (num_steps - len(ask_volumes))
    })
    # For time 0, set bid/ask prices based on the first mid_price
    if num_steps > 0:
        df.loc[0, 'bid_price'] = df.loc[0, 'mid_price'] - base_spread / 2
        df.loc[0, 'ask_price'] = df.loc[0, 'mid_price'] + base_spread / 2
        df.loc[0, 'bid_volume'] = max(1, int(np.random.normal(50, 10)))
        df.loc[0, 'ask_volume'] = max(1, int(np.random.normal(50, 10)))


    logger.debug(f"Generated market data head: {df.head()}")
    return df

def calculate_dynamic_spread(state: dict) -> tuple[float, float]:
    """
    Calculates the dynamic bid-ask spread based on inventory and market parameters.
    This function implements a simplified version of the Avellaneda-Stoikov model for spread calculation.

    Parameters
    ----------
    state : dict
        The current market state dictionary, containing parameters like `risk_aversion`, `volatility`,
        `kappa`, `lambda_val`, `inventory`, and `max_inventory`.

    Returns
    -------
    tuple[float, float]
        A tuple containing the calculated `bid_offset` and `ask_offset` from the current `mid_price`.

    Examples
    --------
    >>> state = create_market_state(10000, 0, 100.0, risk_aversion=0.001, volatility=0.01, kappa=0.01, lambda_val=1.0)
    >>> bid_off, ask_off = calculate_dynamic_spread(state)
    >>> print(f"Bid Offset: {bid_off:.4f}, Ask Offset: {ask_off:.4f}")
    """
    logger.info("Calculating dynamic spread.")

    gamma = state['risk_aversion']
    sigma = state['volatility']
    T = 1.0 # Remaining time horizon (simplified to a constant for steady state)
    q = state['inventory'] # Current inventory
    kappa = state['kappa'] # Order book resilience parameter

    # Optimal spread components based on Avellaneda-Stoikov
    spread_from_vol_risk = gamma * sigma**2 * T
    spread_from_order_flow = (2 / gamma) * math.log(1 + gamma / kappa)
    total_optimal_spread_width = spread_from_vol_risk + spread_from_order_flow

    # Inventory adjustment for reservation price
    reservation_price_offset = -q * gamma * sigma**2 * T

    # Calculate bid and ask offsets from mid_price
    # bid_price = mid_price + reservation_price_offset - 0.5 * total_optimal_spread_width
    # ask_price = mid_price + reservation_price_offset + 0.5 * total_optimal_spread_width
    bid_offset = reservation_price_offset - 0.5 * total_optimal_spread_width
    ask_offset = reservation_price_offset + 0.5 * total_optimal_spread_width

    # Ensure offsets are symmetrical and account for tick size
    tick_size = state['tick_size']

    # Ensure bid_offset is negative and ask_offset is positive relative to mid (after reservation adjustment)
    # And snap to tick size
    bid_offset = round(bid_offset / tick_size) * tick_size
    ask_offset = round(ask_offset / tick_size) * tick_size

    # Ensure at least one tick spread
    if (ask_offset - bid_offset) < tick_size * 2:
        avg_offset = (ask_offset + bid_offset) / 2
        bid_offset = avg_offset - tick_size
        ask_offset = avg_offset + tick_size

    logger.debug(f"Calculated bid_offset: {bid_offset:.4f}, ask_offset: {ask_offset:.4f} for inventory: {q}")
    return bid_offset, ask_offset

def place_orders(state: dict, current_market_mid_price: float) -> tuple[float, float, int]:
    """
    Places bid and ask orders based on dynamic spread calculation.

    Parameters
    ----------
    state : dict
        The current market maker's state dictionary.
    current_market_mid_price : float
        The mid-price observed in the market at the current time step.

    Returns
    -------
    tuple[float, float, int]
        A tuple containing the `bid_price`, `ask_price`, and `order_quantity` for each order.

    Examples
    --------
    >>> state = create_market_state(10000, 0, 100.0)
    >>> bid_p, ask_p, qty = place_orders(state, 100.0)
    >>> print(f"Placed Bid: {bid_p}, Placed Ask: {ask_p}, Quantity: {qty}")
    """
    logger.info(f"Placing orders at market mid_price: {current_market_mid_price:.4f}")
    bid_offset, ask_offset = calculate_dynamic_spread(state)

    # Adjust bid and ask prices based on current market mid-price and calculated offsets
    bid_price = round((current_market_mid_price + bid_offset) / state['tick_size']) * state['tick_size']
    ask_price = round((current_market_mid_price + ask_offset) / state['tick_size']) * state['tick_size']

    order_quantity = 10 # Fixed order quantity for simplicity

    logger.debug(f"Placed orders: Bid={bid_price:.4f}, Ask={ask_price:.4f}, Qty={order_quantity}")
    return bid_price, ask_price, order_quantity

def execute_trade(
    state: dict,
    mm_bid_price: float,
    mm_ask_price: float,
    order_quantity: int,
    market_best_bid: float,
    market_best_ask: float
) -> tuple[dict, int, int]:
    """
    Simulates trade execution against market's best bid/ask and updates state.
    Includes random jitter for timing/backoff mechanisms.

    Parameters
    ----------
    state : dict
        The current market maker's state dictionary.
    mm_bid_price : float
        The bid price placed by the market maker.
    mm_ask_price : float
        The ask price placed by the market maker.
    order_quantity : int
        The quantity of each order.
    market_best_bid : float
        The best bid price currently in the market.
    market_best_ask : float
        The best ask price currently in the market.

    Returns
    -------
    tuple[dict, int, int]
        The updated market maker's state dictionary after potential trades, and the number of shares bought and sold.

    Examples
    --------
    >>> state = create_market_state(10000, 0, 100.0)
    >>> updated_state, bought, sold = execute_trade(state, 99.9, 100.1, 10, 100.05, 100.15)
    >>> print(f"Inventory: {updated_state['inventory']}, Cash: {updated_state['cash']}")
    """
    logger.info(f"Attempting trade execution. MM Bid: {mm_bid_price:.4f}, MM Ask: {mm_ask_price:.4f}")
    shares_bought = 0
    shares_sold = 0

    # Introduce random jitter for simulating latency/backoff
    time.sleep(random.uniform(0.001, 0.005))

    # Check for fill on the market maker's ask (someone buys from us)
    if market_best_bid >= mm_ask_price:
        # Ensure we don't oversell more than we have or beyond min_inventory
        if state['inventory'] - order_quantity >= state['min_inventory']:
            state['cash'] += mm_ask_price * order_quantity
            state['inventory'] -= order_quantity
            shares_sold = order_quantity
            logger.debug(f"Sold {order_quantity} at {mm_ask_price:.4f}. New inventory: {state['inventory']}, cash: {state['cash']:.2f}")
        else:
            logger.warning(f"Attempted to sell {order_quantity} but would exceed min_inventory. Current inventory: {state['inventory']}")

    # Check for fill on the market maker's bid (we buy from someone)
    elif market_best_ask <= mm_bid_price:
        # Ensure we don't overbuy beyond max_inventory
        if state['inventory'] + order_quantity <= state['max_inventory']:
            state['cash'] -= mm_bid_price * order_quantity
            state['inventory'] += order_quantity
            shares_bought = order_quantity
            logger.debug(f"Bought {order_quantity} at {mm_bid_price:.4f}. New inventory: {state['inventory']}, cash: {state['cash']:.2f}")
        else:
            logger.warning(f"Attempted to buy {order_quantity} but would exceed max_inventory. Current inventory: {state['inventory']}")

    return state, shares_bought, shares_sold

def calculate_pnl(state: dict, current_mid_price: float) -> float:
    """
    Calculates the market maker's current Profit and Loss (PnL).

    Parameters
    ----------
    state : dict
        The current market maker's state dictionary.
    current_mid_price : float
        The current mid-price of the asset in the market.

    Returns
    -------
    float
        The current total PnL.

    Examples
    --------
    >>> state = create_market_state(10000, 10, 100.0)
    >>> pnl = calculate_pnl(state, 100.5)
    >>> print(f"Current PnL: {pnl:.2f}")
    """
    logger.info(f"Calculating PnL at mid_price: {current_mid_price:.4f}")
    # PnL = Cash + Inventory * Current Mid Price
    pnl = state['cash'] + (state['inventory'] * current_mid_price)
    logger.debug(f"Current PnL: {pnl:.2f} (Cash: {state['cash']:.2f}, Inventory: {state['inventory']}, Mid: {current_mid_price:.4f})")
    return pnl

def update_metrics(
    metrics_history: dict,
    state: dict,
    current_market_mid_price: float,
    mm_bid_price: float,
    mm_ask_price: float,
    shares_bought: int,
    shares_sold: int
) -> dict:
    """
    Updates a dictionary of historical metrics with the current state and trade information.

    Parameters
    ----------
    metrics_history : dict
        A dictionary to store historical metrics.
    state : dict
        The current market maker's state dictionary.
    current_market_mid_price : float
        The mid-price from the market simulation at the current step.
    mm_bid_price : float
        The market maker's bid price at the current step.
    mm_ask_price : float
        The market maker's ask price at the current step.
    shares_bought : int
        Shares bought in the current step.
    shares_sold : int
        Shares sold in the current step.

    Returns
    -------
    dict
        The updated `metrics_history` dictionary.

    Examples
    --------
    >>> history = collections.defaultdict(list)
    >>> state = create_market_state(10000, 0, 100.0)
    >>> updated_history = update_metrics(history, state, 100.0, 99.9, 100.1, 0, 0)
    >>> print(updated_history['pnl'][-1])
    """
    logger.info("Updating metrics history.")
    current_pnl = calculate_pnl(state, current_market_mid_price)
    spread = mm_ask_price - mm_bid_price

    metrics_history['time'].append(state['time'])
    metrics_history['mid_price'].append(current_market_mid_price)
    metrics_history['mm_bid_price'].append(mm_bid_price)
    metrics_history['mm_ask_price'].append(mm_ask_price)
    metrics_history['mm_spread'].append(spread)
    metrics_history['inventory'].append(state['inventory'])
    metrics_history['cash'].append(state['cash'])
    metrics_history['pnl'].append(current_pnl)
    metrics_history['shares_bought'].append(shares_bought)
    metrics_history['shares_sold'].append(shares_sold)
    logger.debug(f"Metrics updated: PnL={current_pnl:.2f}, Inventory={state['inventory']}, Spread={spread:.4f}")
    return metrics_history

def run_simulation(initial_state: dict, market_data_df: pd.DataFrame) -> pd.DataFrame:
    """
    Runs the market making simulation over the provided market data.

    Parameters
    ----------
    initial_state : dict
        The initial market maker's state.
    market_data_df : pd.DataFrame
        A DataFrame containing simulated market data for the simulation period.

    Returns
    -------
    pd.DataFrame
        A DataFrame containing the full history of recorded metrics during the simulation.

    Examples
    --------
    >>> state = create_market_state(10000, 0, 100.0)
    >>> market_df = generate_market_data(state, 100)
    >>> simulation_results = run_simulation(state, market_df)
    >>> print(simulation_results.head())
    """
    logger.info("Starting market making simulation.")
    current_state = initial_state.copy()
    metrics_history = collections.defaultdict(list)

    for index, row in tqdm(market_data_df.iterrows(), total=len(market_data_df), desc="Simulating Market Making"):
        current_state['time'] = row['time']
        current_state['mid_price'] = row['mid_price']

        mm_bid_price, mm_ask_price, order_quantity = place_orders(current_state, row['mid_price'])

        # Execute trades against the simulated market's best bid/ask
        current_state, shares_bought, shares_sold = execute_trade(
            current_state,
            mm_bid_price,
            mm_ask_price,
            order_quantity,
            row['bid_price'], # Use market's bid for our ask fill
            row['ask_price']  # Use market's ask for our bid fill
        )

        metrics_history = update_metrics(
            metrics_history,
            current_state,
            row['mid_price'],
            mm_bid_price,
            mm_ask_price,
            shares_bought,
            shares_sold
        )
    logger.info("Market making simulation completed.")
    return pd.DataFrame(metrics_history)

# --- Simulation Parameters ---
SIM_STEPS = 500
INITIAL_CASH = 100000.0
INITIAL_INVENTORY = 0
INITIAL_MID_PRICE = 100.0
RISK_AVERSION = 0.5 # Increased risk aversion for more sensitive spreads
VOLATILITY = 0.02 # Increased market volatility
KAPPA = 10.0 # Increased order book resilience parameter
LAMBDA_VAL = 1.0 # Order arrival rate intensity (kept similar or slightly adjusted for balance)

logger.info("Setting up initial market state and generating market data.")
# Initialize market state
initial_mm_state = create_market_state(
    initial_cash=INITIAL_CASH,
    initial_inventory=INITIAL_INVENTORY,
    initial_mid_price=INITIAL_MID_PRICE,
    risk_aversion=RISK_AVERSION,
    volatility=VOLATILITY,
    kappa=KAPPA,
    lambda_val=LAMBDA_VAL
)

# Generate simulated market data
simulated_market_data = generate_market_data(initial_mm_state, SIM_STEPS)

# Run the simulation
simulation_results_df = run_simulation(initial_mm_state, simulated_market_data)

logger.info("Simulation complete. Displaying summary statistics.")
# Display summary statistics of the simulation results
print("\n--- Simulation Results Summary ---")
display(simulation_results_df.describe())

plt.figure(figsize=(15, 8))
sns.set_style("whitegrid")

# Plot mid-price
sns.lineplot(x='time', y='mid_price', data=simulation_results_df, label='Market Mid-Price', color='black', alpha=0.7)

# Plot market maker's bid and ask prices
sns.lineplot(x='time', y='mm_bid_price', data=simulation_results_df, label='MM Bid Price', color='red', linestyle='--')
sns.lineplot(x='time', y='mm_ask_price', data=simulation_results_df, label='MM Ask Price', color='green', linestyle='--')

plt.title('Market Mid-Price vs. Market Maker Bid/Ask Prices')
plt.xlabel('Time Step')
plt.ylabel('Price')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()

fig, axes = plt.subplots(2, 1, figsize=(15, 10), sharex=True)
sns.set_style("whitegrid")

# Plot Inventory
sns.lineplot(x='time', y='inventory', data=simulation_results_df, ax=axes[0], color='purple')
axes[0].axhline(y=0, color='gray', linestyle=':', linewidth=0.8, label='Zero Inventory')
axes[0].set_title('Market Maker Inventory Over Time')
axes[0].set_ylabel('Inventory (Shares)')
axes[0].legend()

# Plot PnL
sns.lineplot(x='time', y='pnl', data=simulation_results_df, ax=axes[1], color='orange')
axes[1].axhline(y=INITIAL_CASH, color='gray', linestyle=':', linewidth=0.8, label='Initial Capital')
axes[1].set_title('Market Maker PnL Over Time')
axes[1].set_xlabel('Time Step')
axes[1].set_ylabel('PnL (USD)')
axes[1].legend()

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(2, 1, figsize=(15, 10), sharex=True)
sns.set_style("whitegrid")

# Plot Dynamic Spread
sns.lineplot(x='time', y='mm_spread', data=simulation_results_df, ax=axes[0], color='teal')
axes[0].set_title('Market Maker Dynamic Bid-Ask Spread')
axes[0].set_ylabel('Spread (USD)')

# Plot Trade Activity (Bought/Sold)
axes[1].fill_between(simulation_results_df['time'], 0, simulation_results_df['shares_bought'], color='blue', alpha=0.3, label='Shares Bought')
axes[1].fill_between(simulation_results_df['time'], 0, -simulation_results_df['shares_sold'], color='red', alpha=0.3, label='Shares Sold')
axes[1].set_title('Trade Activity (Shares Bought/Sold per Step)')
axes[1].set_xlabel('Time Step')
axes[1].set_ylabel('Shares')
axes[1].legend()

plt.tight_layout()
plt.show()