import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker

def partial_fill_simulator(
    order_quantity: float,
    order_book_side: pd.DataFrame,
    order_type: str
) -> dict:
    """
    Simulates the partial filling of a market order against a given order book side.

    This function processes an order against the provided order book, calculating
    how much is filled, what remains, and the average execution price.

    Args:
        order_quantity (float): The total quantity of the asset to buy or sell.
        order_book_side (pd.DataFrame): A DataFrame representing one side of the
                                         order book (e.g., 'asks' for a buy order).
                                         It must have 'price' and 'quantity' columns.
                                         Prices should be sorted appropriately (ascending for asks,
                                         descending for bids).
        order_type (str): The type of order, 'buy' or 'sell'. This determines
                          which side of the order book is consumed.

    Returns:
        dict: A dictionary containing:
            'filled_quantity': The total quantity filled.
            'remaining_quantity': The quantity of the order that could not be filled.
            'average_execution_price': The volume-weighted average price of the filled quantity.
            'filled_details': A list of dictionaries, each containing 'price', 'quantity_filled'
                              for each level where the order was partially filled.
            'initial_best_price': The best price available at the start of the simulation.
    """
    if order_book_side.empty:
        return {
            'filled_quantity': 0.0,
            'remaining_quantity': order_quantity,
            'average_execution_price': None,
            'filled_details': [],
            'initial_best_price': None
        }

    filled_quantity = 0.0
    total_cost_or_revenue = 0.0
    filled_details = []
    initial_best_price = order_book_side['price'].iloc[0]

    remaining_order_quantity = order_quantity

    for _, row in order_book_side.iterrows():
        price = row['price']
        available_quantity = row['quantity']

        if remaining_order_quantity <= 0:
            break

        quantity_to_fill = min(remaining_order_quantity, available_quantity)

        filled_quantity += quantity_to_fill
        total_cost_or_revenue += quantity_to_fill * price
        remaining_order_quantity -= quantity_to_fill

        filled_details.append({
            'price': price,
            'quantity_filled': quantity_to_fill
        })

    average_execution_price = total_cost_or_revenue / filled_quantity if filled_quantity > 0 else None

    return {
        'filled_quantity': filled_quantity,
        'remaining_quantity': remaining_order_quantity,
        'average_execution_price': average_execution_price,
        'filled_details': filled_details,
        'initial_best_price': initial_best_price
    }

print("Partial fill simulator function defined.")

# Seed for reproducibility
np.random.seed(42)

# Define parameters for the mock order book
best_ask_price = 100.00
num_price_levels = 20
price_increment = 0.01 # Smallest tick size

# Generate ask prices (increasing from best ask)
ask_prices = np.linspace(best_ask_price, best_ask_price + (num_price_levels - 1) * price_increment, num_price_levels)

# Generate quantities for each ask price level
# Quantities can be somewhat random, but often decrease further from the best price
ask_quantities = np.random.randint(50, 500, size=num_price_levels)
# Let's make the quantities generally larger closer to the best price, then drop off
ask_quantities = ask_quantities * (1 + np.exp(-(np.arange(num_price_levels) / 5)))
ask_quantities = np.round(ask_quantities, 0)

mock_asks_df = pd.DataFrame({
    'price': ask_prices,
    'quantity': ask_quantities
})

print("Mock Ask Order Book (first 5 levels):")
display(mock_asks_df.head())

print("\nMock Ask Order Book (last 5 levels):")
display(mock_asks_df.tail())

buy_order_quantity = 2500  # A relatively large buy order

print(f"Simulating a BUY order of {buy_order_quantity} units.\n")

simulation_results = partial_fill_simulator(
    order_quantity=buy_order_quantity,
    order_book_side=mock_asks_df,
    order_type='buy'
)

print("--- Simulation Results ---")
print(f"Total Order Quantity: {buy_order_quantity}")
print(f"Filled Quantity: {simulation_results['filled_quantity']:.2f}")
print(f"Remaining Quantity: {simulation_results['remaining_quantity']:.2f}")
if simulation_results['average_execution_price'] is not None:
    print(f"Average Execution Price: {simulation_results['average_execution_price']:.4f}")
    print(f"Initial Best Ask Price: {simulation_results['initial_best_price']:.4f}")
    slippage = simulation_results['average_execution_price'] - simulation_results['initial_best_price']
    print(f"Slippage: {slippage:.4f}")
else:
    print("Could not calculate average execution price (no fill).")

print("\n--- Filled Details ---")
filled_df = pd.DataFrame(simulation_results['filled_details'])
display(filled_df)

fig, ax = plt.subplots(figsize=(12, 7))

# Prepare data for plotting
plot_df = mock_asks_df.copy()
plot_df['quantity_filled'] = 0.0

for detail in simulation_results['filled_details']:
    price = detail['price']
    qty_filled = detail['quantity_filled']
    plot_df.loc[plot_df['price'] == price, 'quantity_filled'] = qty_filled

plot_df['quantity_remaining_at_level'] = plot_df['quantity'] - plot_df['quantity_filled']

# Create stacked bar chart
bar_width = price_increment * 0.8 # Make bars slightly narrower than price increment

ax.bar(plot_df['price'], plot_df['quantity_filled'], label='Quantity Filled', color='skyblue', width=bar_width)
ax.bar(plot_df['price'], plot_df['quantity_remaining_at_level'], bottom=plot_df['quantity_filled'], label='Remaining Liquidity at Level', color='lightgray', width=bar_width)

# Add a line for the average execution price
if simulation_results['average_execution_price'] is not None:
    ax.axvline(x=simulation_results['average_execution_price'], color='red', linestyle='--',
               label=f"Avg. Exec. Price: {simulation_results['average_execution_price']:.2f}")
    ax.text(simulation_results['average_execution_price'] + 0.005, ax.get_ylim()[1] * 0.9,
            'Avg. Exec. Price',
            color='red', rotation=90, va='top')

# Add a line for the initial best price
ax.axvline(x=simulation_results['initial_best_price'], color='green', linestyle='--',
           label=f"Initial Best Ask: {simulation_results['initial_best_price']:.2f}")
ax.text(simulation_results['initial_best_price'] - 0.005, ax.get_ylim()[1] * 0.9,
        'Initial Best Ask',
        color='green', rotation=90, va='top', ha='right')

ax.set_title(f'Order Book Liquidity Consumption for Buy Order of {buy_order_quantity} units')
ax.set_xlabel('Price')
ax.set_ylabel('Quantity')
ax.legend()
ax.grid(axis='y', linestyle='--', alpha=0.7)

# Format x-axis to show more price precision
formatter = mticker.FormatStrFormatter('$%.2f')
ax.xaxis.set_major_formatter(formatter)
plt.xticks(rotation=45, ha='right')
plt.tight_layout()
plt.show()

order_quantities = np.arange(100, 5000, 200) # Test various order quantities
slippages = []
avg_prices = []

for qty in order_quantities:
    results = partial_fill_simulator(
        order_quantity=qty,
        order_book_side=mock_asks_df,
        order_type='buy'
    )
    if results['average_execution_price'] is not None:
        slippage = results['average_execution_price'] - results['initial_best_price']
        slippages.append(slippage)
        avg_prices.append(results['average_execution_price'])
    else:
        slippages.append(0.0) # Or NaN if preferred for unfilled orders
        avg_prices.append(mock_asks_df['price'].iloc[0]) # Initial best price if no fill

fig, ax1 = plt.subplots(figsize=(12, 7))

color = 'tab:blue'
ax1.set_xlabel('Order Quantity')
ax1.set_ylabel('Slippage (Price difference)', color=color)
ax1.plot(order_quantities, slippages, color=color, marker='o', linestyle='-')
ax1.tick_params(axis='y', labelcolor=color)
ax1.set_title('Slippage and Average Execution Price vs. Order Quantity')
ax1.grid(True, linestyle='--', alpha=0.6)

# Create a second y-axis for average execution price
ax2 = ax1.twinx()
color = 'tab:red'
ax2.set_ylabel('Average Execution Price', color=color)
ax2.plot(order_quantities, avg_prices, color=color, marker='x', linestyle='--')
ax2.tick_params(axis='y', labelcolor=color)

# Format y-axes to show currency
formatter = mticker.FormatStrFormatter('$%.4f')
ax1.yaxis.set_major_formatter(formatter)
ax2.yaxis.set_major_formatter(formatter)

# Add horizontal line for the initial best price
ax1.axhline(0, color='gray', linestyle=':', linewidth=0.8, label='No Slippage (Initial Best Price)')
ax2.axhline(mock_asks_df['price'].iloc[0], color='gray', linestyle=':', linewidth=0.8)

fig.tight_layout()  # otherwise the right y-label is slightly clipped
plt.show()