import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import warnings

# Suppress all warnings for cleaner output in a demonstration context.
warnings.filterwarnings("ignore")

def calculate_spread_price(mid_price: float, spread_bps: float) -> float:
    """Return the half-spread offset in price units.

    Parameters
    ----------
    mid_price : Reference mid-market price.
    spread_bps: Bid-ask spread in basis points.

    Returns
    -------
    float: Calculated half-spread value.
    """
    # Calculate spread based on mid-price and basis points, then halve it.
    return mid_price * (spread_bps / 10_000) / 2

def get_best_ask(mid_price: float, spread_price_val: float) -> float:
    """Return the best ask price.

    Parameters
    ----------
    mid_price       : Reference mid-market price.
    spread_price_val: Half-spread value.

    Returns
    -------
    float: Best ask price.
    """
    # Best ask is mid-price plus the half-spread.
    return mid_price + spread_price_val

def get_best_bid(mid_price: float, spread_price_val: float) -> float:
    """Return the best bid price.

    Parameters
    ----------
    mid_price       : Reference mid-market price.
    spread_price_val: Half-spread value.

    Returns
    -------
    float: Best bid price.
    """
    # Best bid is mid-price minus the half-spread.
    return mid_price - spread_price_val

def get_ask_levels(mid_price: float, best_ask: float, depth_levels: int, size_per_level: float) -> list[tuple[float, float]]:
    """Return ask side: list of (price, size) tuples, ascending.

    Parameters
    ----------
    mid_price     : Reference mid-market price (for tick calculation).
    best_ask      : The best ask price.
    depth_levels  : Number of price levels to generate.
    size_per_level: Quantity available at each level.

    Returns
    -------
    list[tuple[float, float]]: List of (price, size) for ask levels.
    """
    # Define a tick as 0.01% of the mid-price for granular price levels.
    tick = mid_price * 0.0001
    # Generate ask levels starting from the best ask, increasing by tick.
    return [(best_ask + i * tick, size_per_level)
            for i in range(depth_levels)]

def get_bid_levels(mid_price: float, best_bid: float, depth_levels: int, size_per_level: float) -> list[tuple[float, float]]:
    """Return bid side: list of (price, size) tuples, descending.

    Parameters
    ----------
    mid_price     : Reference mid-market price (for tick calculation).
    best_bid      : The best bid price.
    depth_levels  : Number of price levels to generate.
    size_per_level: Quantity available at each level.

    Returns
    -------
    list[tuple[float, float]]: List of (price, size) for bid levels.
    """
    # Define a tick as 0.01% of the mid-price for granular price levels.
    tick = mid_price * 0.0001
    # Generate bid levels starting from the best bid, decreasing by tick.
    return [(best_bid - i * tick, size_per_level)
            for i in range(depth_levels)]

def simulate_market_fill(mid_price: float, spread_bps: float, depth_levels: int, size_per_level: float, quantity: float, side: str) -> dict:
    """
    Simulate market order execution by walking the order book.

    Parameters
    ----------
    mid_price      : Reference mid-market price.
    spread_bps     : Bid-ask spread in basis points.
    depth_levels   : Number of price levels on each side.
    size_per_level : Quantity available at each price level.
    quantity       : Base asset quantity to fill.
    side           : 'buy' | 'sell'

    Returns
    -------
    dict: Contains avg_fill_price, slippage_bps, total_cost, unfilled_qty.
    """
    # Calculate spread and best bid/ask prices
    spread_price_val = calculate_spread_price(mid_price, spread_bps)
    best_ask_val = get_best_ask(mid_price, spread_price_val)
    best_bid_val = get_best_bid(mid_price, spread_price_val)

    # Determine the relevant side of the order book and reference price
    if side == "buy":
        levels = get_ask_levels(mid_price, best_ask_val, depth_levels, size_per_level)
        ref_price = best_ask_val
    else: # side == "sell"
        levels = get_bid_levels(mid_price, best_bid_val, depth_levels, size_per_level)
        ref_price = best_bid_val

    remaining_qty = quantity
    total_cost_incurred = 0.0
    total_filled_qty = 0.0

    # Iterate through order book levels to simulate fills
    for price, available_size in levels:
        fill_amount = min(remaining_qty, available_size)
        total_cost_incurred   += fill_amount * price
        total_filled_qty += fill_amount
        remaining_qty    -= fill_amount
        if remaining_qty <= 0:
            break

    # Handle cases where no fill occurs
    if total_filled_qty == 0:
        return {"avg_fill_price": 0, "slippage_bps": 0, "total_cost": 0, "unfilled_qty": quantity}

    # Calculate average fill price and slippage
    avg_price    = total_cost_incurred / total_filled_qty
    slippage_bps = abs((avg_price - ref_price) / ref_price) * 10_000

    return {
        "avg_fill_price": round(avg_price, 2),
        "slippage_bps":   round(slippage_bps, 4),
        "total_cost":     round(total_cost_incurred, 2),
        "unfilled_qty":   round(remaining_qty, 6),
    }

def simulate_limit_fill(mid_price: float, spread_bps: float, quantity: float, side: str,
                         limit_price: float, ticks_offset: int = 3) -> dict:
    """
    Simulate limit order fill probability based on price offset.

    A limit order fills only if the market price reaches the limit price.
    Fill probability decays as the limit is farther from the best price.

    Parameters
    ----------
    mid_price   : Reference mid-market price.
    spread_bps  : Bid-ask spread in basis points.
    quantity    : Quantity to fill.
    side        : 'buy' | 'sell'
    limit_price : The specified limit price.
    ticks_offset: Not directly used in current probability model, but signifies offset intent.

    Returns
    -------
    dict: Contains fill_probability, avg_fill_price, slippage_bps, maker_rebate_usd.
    """
    # Define tick size and calculate best bid/ask prices
    tick = mid_price * 0.0001
    spread_price_val = calculate_spread_price(mid_price, spread_bps)
    best_ask_val = get_best_ask(mid_price, spread_price_val)
    best_bid_val = get_best_bid(mid_price, spread_price_val)

    # Set reference price based on order side
    ref_price = best_ask_val if side == "buy" else best_bid_val

    # Calculate distance in ticks from the reference price
    distance_ticks = abs(limit_price - ref_price) / tick

    # Determine fill probability:
    #   1.0 if aggressive (crossing the book).
    #   Decays linearly with distance for passive orders.
    is_crossing   = (side == "buy" and limit_price >= ref_price) or \
                    (side == "sell" and limit_price <= ref_price)
    fill_prob     = 1.0 if is_crossing else max(0, 1 - 0.08 * distance_ticks)

    # Limit orders, if filled, guarantee the limit price, thus zero slippage.
    slippage_bps  = 0.0
    # Maker rebate is applied for passive orders (not crossing the book).
    maker_rebate  = 0.0 if is_crossing else quantity * limit_price * 0.0001   # Example: 0.01% rebate

    return {
        "fill_probability": round(fill_prob, 4),
        "avg_fill_price":   round(limit_price, 2),
        "slippage_bps":     slippage_bps,
        "maker_rebate_usd": round(maker_rebate, 4),
    }

def compute_order_fees(notional_usd: float, order_type: str,
                        exchange: str = "binance") -> dict:
    """
    Calculate trading fees for a given order type and exchange.

    Parameters
    ----------
    notional_usd : Total order value in USD.
    order_type   : 'market' (taker) | 'limit' (maker).
    exchange     : Exchange name for fee schedule lookup.

    Returns
    -------
    dict: Contains fee_usd and fee_bps.
    """
    # Representative fee schedules (e.g., VIP 0 tier, USDT perpetuals).
    # Fee rates are expressed as decimals (e.g., 0.0004 for 0.04%).
    fee_schedules = {
        "binance": {"taker": 0.0004, "maker": 0.0002},
        "bybit":   {"taker": 0.0006, "maker": 0.0001},
        "okx":     {"taker": 0.0005, "maker": 0.0002},
        "kraken":  {"taker": 0.0005, "maker": 0.0002},
    }
    # Retrieve schedule for the specified exchange, defaulting if not found.
    schedule  = fee_schedules.get(exchange.lower(), {"taker": 0.0005, "maker": 0.0002})
    # Apply taker rate for market orders, maker rate for limit orders.
    fee_rate  = schedule["taker"] if order_type == "market" else schedule["maker"]
    fee_usd   = notional_usd * fee_rate

    return {
        "fee_usd": round(fee_usd, 4),
        "fee_bps": round(fee_rate * 10_000, 2),
    }

# Define a constant notional value for fee calculation across exchanges.
notional = 10_000   # Notional value in USD for a hypothetical trade.

results = []
# Iterate through defined exchanges and order types to compute fees.
for ex in ["binance", "bybit", "okx", "kraken"]:
    for otype in ["market", "limit"]:
        fees = compute_order_fees(notional, otype, ex)
        results.append({"Exchange": ex.capitalize(), "Order Type": otype.capitalize(),
                        "Fee (USD)": fees["fee_usd"], "Fee (bps)": fees["fee_bps"]})

# Convert results to a pandas DataFrame for structured display.
df_fees = pd.DataFrame(results)

# Display fee comparison results in a formal table format.
print(f"Fee Comparison — ${notional:,} Notional")
display(df_fees)

# Define simulation parameters for market orders, establishing the synthetic market environment.
mid_price_val = 50_000         # Reference mid-market price for the asset.
spread_bps_val = 2             # Bid-ask spread in basis points.
depth_levels_val = 20          # Number of price levels simulated on each side of the book.
size_per_level_val = 0.2       # Quantity (e.g., BTC) available at each price level.

# Define a range of market order sizes (in BTC) to be simulated.
order_sizes  = [0.1, 0.5, 1.0, 2.0, 5.0, 10.0]
market_fills_data = [] # Initialize list to store results for each order size.

# Execute market order simulations for each defined order size.
for qty in order_sizes:
    # Simulate market fill for a 'buy' order.
    result = simulate_market_fill(mid_price_val, spread_bps_val, depth_levels_val, size_per_level_val, qty, "buy")
    # Compute associated taker fees for the order's notional value.
    fees   = compute_order_fees(qty * mid_price_val, "market", "binance")
    market_fills_data.append({
        "Order Size (BTC)":    qty,
        "Avg Fill Price":      result["avg_fill_price"],
        "Slippage (bps)":      result["slippage_bps"],
        "Unfilled Qty (BTC)":  result["unfilled_qty"],
        "Fee (USD)":           fees["fee_usd"],
        "Total Cost (USD)":    result["total_cost"] + fees["fee_usd"], # Total cost includes fill cost and fees.
    })

# Convert the simulation results to a pandas DataFrame.
df_market = pd.DataFrame(market_fills_data)

# Display market order simulation results in a formal table format.
print("Market Order Simulation Results")
display(df_market)

# Define simulation parameters for limit orders, consistent with market order simulation for comparability.
mid_price_val = 50_000 # Reference mid-market price.
spread_bps_val = 2     # Bid-ask spread in basis points.

# Calculate the granular tick size based on the mid-price.
tick         = mid_price_val * 0.0001
# Define a range of price offsets (in ticks) from the best ask to evaluate fill probability.
offset_ticks = list(range(0, 15))
limit_results_data = [] # Initialize list to store results for each offset.

# Calculate the best ask price, which serves as the reference for limit order placement.
spread_price_val = calculate_spread_price(mid_price_val, spread_bps_val)
best_ask_val = get_best_ask(mid_price_val, spread_price_val)

# Execute limit order simulations for each offset from the best ask.
for ticks in offset_ticks:
    # Calculate the limit price for a passive 'buy' order (below the best ask).
    limit_price = best_ask_val - ticks * tick
    # Simulate the limit order fill characteristics.
    result = simulate_limit_fill(mid_price_val, spread_bps_val, 1.0, "buy", limit_price, ticks)
    limit_results_data.append({
        "Ticks from Ask": ticks,
        "Limit Price":    round(limit_price, 2),
        "Fill Prob (%)":  result["fill_probability"] * 100, # Convert probability to percentage.
        "Slippage (bps)": result["slippage_bps"],
        "Maker Rebate $": result["maker_rebate_usd"],
    })

# Convert the simulation results to a pandas DataFrame.
df_limit = pd.DataFrame(limit_results_data)

# Display limit order simulation results in a formal table format.
print("Limit Order Fill Probability vs Price Offset")
display(df_limit)

fig, axes = plt.subplots(1, 3, figsize=(18, 6)) # Create a figure with three subplots for comprehensive visualization.
fig.suptitle("Market vs Limit Order Analysis", fontsize=16, fontweight="bold", y=1.02) # Set a main title for the figure.

# Plot 1: Market Order Slippage vs. Order Size
# This plot illustrates how slippage (in basis points) increases with larger market order sizes.
axes[0].plot(df_market["Order Size (BTC)"], df_market["Slippage (bps)"],
             marker="o", color="#e74c3c", linewidth=2, linestyle='-')
axes[0].set_xlabel("Order Size (BTC)")
axes[0].set_ylabel("Slippage (bps)")
axes[0].set_title("Market Order Slippage vs. Size")
axes[0].grid(True, linestyle='--', alpha=0.6)

# Plot 2: Limit Order Fill Probability vs. Price Offset
# This bar chart shows the decreasing fill probability of a limit order as it is placed further from the best ask price.
axes[1].bar(df_limit["Ticks from Ask"], df_limit["Fill Prob (%)"],
            color="#3498db", alpha=0.8, edgecolor="black", linewidth=0.7)
axes[1].set_xlabel("Ticks from Best Ask")
axes[1].set_ylabel("Fill Probability (%)")
axes[1].set_title("Limit Order Fill Probability vs. Offset")
axes[1].grid(True, axis="y", linestyle='--', alpha=0.6)

# Plot 3: Transaction Fee Comparison by Exchange and Order Type
# This bar chart compares transaction fees (in USD) for market and limit orders across different exchanges.
pivot = df_fees.pivot(index="Exchange", columns="Order Type", values="Fee (USD)")
pivot.plot(kind="bar", ax=axes[2], color=["#e74c3c", "#2ecc71"], # Red for Market, Green for Limit.
           edgecolor="black", linewidth=0.7)
axes[2].set_title(f"Transaction Fee Comparison — ${notional:,} Notional")
axes[2].set_ylabel("Fee (USD)")
axes[2].set_xlabel("") # Remove x-axis label as 'Exchange' is self-explanatory.
axes[2].legend(title="Order Type")
axes[2].tick_params(axis="x", rotation=45) # Rotate x-axis labels for improved readability.
axes[2].grid(True, axis="y", linestyle='--', alpha=0.6)

plt.tight_layout() # Adjust subplot parameters for a tight layout.
plt.savefig("order_type_comparison.png", dpi=300, bbox_inches="tight") # Save the figure with high resolution.
plt.show() # Display the generated plots.