import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import random

def generate_mock_order_book(mid_price: float = 100.0, num_levels: int = 10, max_qty: int = 100) -> dict:
    """
    Generates mock order book data around a given mid-price.

    Args:
        mid_price (float): The central price around which bids and asks are generated.
        num_levels (int): The number of price levels to generate for bids and asks.
        max_qty (int): The maximum quantity for any single order at a price level.

    Returns:
        dict: A dictionary containing 'bids' and 'asks', where each is a list of (price, quantity) tuples.
    """
    bids = []
    asks = []

    # Generate bids (prices decreasing from mid_price)
    for i in range(num_levels):
        price = round(mid_price - (0.01 * (i + 1)) - random.uniform(0, 0.005), 4) # Decreasing prices
        qty = random.randint(10, max_qty)
        bids.append((price, qty))
    bids.sort(key=lambda x: x[0], reverse=True) # Ensure bids are sorted descending by price

    # Generate asks (prices increasing from mid_price)
    for i in range(num_levels):
        price = round(mid_price + (0.01 * (i + 1)) + random.uniform(0, 0.005), 4) # Increasing prices
        qty = random.randint(10, max_qty)
        asks.append((price, qty))
    asks.sort(key=lambda x: x[0]) # Ensure asks are sorted ascending by price

    return {'bids': bids, 'asks': asks}

# Generate mock data
mock_order_book_data = generate_mock_order_book(mid_price=100.0, num_levels=15, max_qty=200)
print("Mock Bid Orders (Price, Quantity):\n", mock_order_book_data['bids'])
print("\nMock Ask Orders (Price, Quantity):\n", mock_order_book_data['asks'])

def create_full_depth_snapshot(order_book_data: dict) -> pd.DataFrame:
    """
    Creates a full depth order book snapshot from raw bid/ask data.

    Args:
        order_book_data (dict): A dictionary with 'bids' and 'asks' lists of (price, quantity) tuples.

    Returns:
        pd.DataFrame: A DataFrame representing the full depth order book snapshot.
            Columns: 'Price', 'Bid_Quantity', 'Ask_Quantity'.
    """
    # Prepare bids DataFrame
    bids_df = pd.DataFrame(order_book_data['bids'], columns=['Price', 'Bid_Quantity'])
    bids_df['Ask_Quantity'] = np.nan # No ask quantity for bid prices

    # Prepare asks DataFrame
    asks_df = pd.DataFrame(order_book_data['asks'], columns=['Price', 'Ask_Quantity'])
    asks_df['Bid_Quantity'] = np.nan # No bid quantity for ask prices

    # Combine and sort by price
    full_df = pd.concat([bids_df, asks_df]).sort_values(by='Price', ascending=False).reset_index(drop=True)

    # Fill NaN values to show the complete picture at each price level
    # This assumes that if a price exists as a bid, it won't be an ask, and vice-versa
    # In a real order book, we would group by price and sum quantities
    full_df = full_df.groupby('Price').agg({
        'Bid_Quantity': 'sum',
        'Ask_Quantity': 'sum'
    }).reset_index().replace(0, np.nan) # Replace 0 sums with NaN if they were initially NaN

    # Sort again to ensure correct order and fill NaNs with 0 for display
    full_df = full_df.sort_values(by='Price', ascending=False).fillna(0).reset_index(drop=True)

    return full_df

# Create and display full depth snapshot
full_depth_snapshot = create_full_depth_snapshot(mock_order_book_data)
print("Full Depth Order Book Snapshot:")
display(full_depth_snapshot.head(10)) # Display top 10 levels for brevity
print("... and bottom 10 levels:")
display(full_depth_snapshot.tail(10))

def plot_full_depth_order_book(snapshot_df: pd.DataFrame, title: str = "Full Depth Order Book Snapshot"):
    """
    Plots the full depth order book snapshot using a bar chart.

    Args:
        snapshot_df (pd.DataFrame): DataFrame with 'Price', 'Bid_Quantity', 'Ask_Quantity' columns.
        title (str): Title for the plot.
    """
    fig, ax = plt.subplots(figsize=(12, 7))

    # Filter out zero quantities for plotting clarity
    bids_to_plot = snapshot_df[snapshot_df['Bid_Quantity'] > 0]
    asks_to_plot = snapshot_df[snapshot_df['Ask_Quantity'] > 0]

    # Plot bids as negative quantities to show on the left
    ax.barh(bids_to_plot['Price'], -bids_to_plot['Bid_Quantity'], color='green', label='Bids', alpha=0.7)
    # Plot asks as positive quantities on the right
    ax.barh(asks_to_plot['Price'], asks_to_plot['Ask_Quantity'], color='red', label='Asks', alpha=0.7)

    ax.set_title(title)
    ax.set_xlabel('Quantity')
    ax.set_ylabel('Price')
    ax.legend()
    ax.grid(True, linestyle='--', alpha=0.6)

    # Format x-axis to show positive values for both bids and asks with a clear demarcation
    ticks = ax.get_xticks()
    ax.set_xticklabels([f'{abs(tick):.0f}' for tick in ticks])

    # Highlight best bid and ask (the prices closest to each other)
    best_bid_price = bids_to_plot['Price'].max() if not bids_to_plot.empty else None
    best_ask_price = asks_to_plot['Price'].min() if not asks_to_plot.empty else None

    if best_bid_price:
        ax.axhline(best_bid_price, color='darkgreen', linestyle='--', linewidth=1, label=f'Best Bid: {best_bid_price}')
    if best_ask_price:
        ax.axhline(best_ask_price, color='darkred', linestyle='--', linewidth=1, label=f'Best Ask: {best_ask_price}')

    # Add a legend for the best bid/ask lines if they exist
    handles, labels = ax.get_legend_handles_labels()
    if best_bid_price and f'Best Bid: {best_bid_price}' not in labels:
        handles.append(plt.Line2D([0, 0], [0, 0], color='darkgreen', linestyle='--', linewidth=1))
        labels.append(f'Best Bid: {best_bid_price}')
    if best_ask_price and f'Best Ask: {best_ask_price}' not in labels:
        handles.append(plt.Line2D([0, 0], [0, 0], color='darkred', linestyle='--', linewidth=1))
        labels.append(f'Best Ask: {best_ask_price}')
    ax.legend(handles, labels)

    plt.tight_layout()
    plt.show()

plot_full_depth_order_book(full_depth_snapshot)

def create_aggregated_snapshot(order_book_data: dict, num_levels: int = 5) -> pd.DataFrame:
    """
    Creates an aggregated order book snapshot, typically showing the top N levels.

    Args:
        order_book_data (dict): A dictionary with 'bids' and 'asks' lists of (price, quantity) tuples.
        num_levels (int): The number of top bid and ask levels to include in the snapshot.

    Returns:
        pd.DataFrame: A DataFrame representing the aggregated order book snapshot.
            Columns: 'Price', 'Bid_Quantity', 'Ask_Quantity'.
    """
    # Sort bids descending by price and take top N
    sorted_bids = sorted(order_book_data['bids'], key=lambda x: x[0], reverse=True)[:num_levels]
    bids_df = pd.DataFrame(sorted_bids, columns=['Price', 'Bid_Quantity'])

    # Sort asks ascending by price and take top N
    sorted_asks = sorted(order_book_data['asks'], key=lambda x: x[0])[:num_levels]
    asks_df = pd.DataFrame(sorted_asks, columns=['Price', 'Ask_Quantity'])

    # Combine dataframes
    # Using outer merge to ensure all prices from both sides are included
    aggregated_df = pd.merge(bids_df, asks_df, on='Price', how='outer').fillna(0)

    # Ensure correct sorting for presentation (highest bid to lowest ask)
    aggregated_df = aggregated_df.sort_values(by='Price', ascending=False).reset_index(drop=True)

    return aggregated_df

# Create and display aggregated snapshot (top 5 levels)
aggregated_snapshot = create_aggregated_snapshot(mock_order_book_data, num_levels=5)
print(f"Aggregated Order Book Snapshot (Top 5 Levels):")
display(aggregated_snapshot)

def plot_cumulative_order_book(snapshot_df: pd.DataFrame, title: str = "Cumulative Order Book Volume"):
    """
    Plots the cumulative bid and ask volumes from an order book snapshot.

    Args:
        snapshot_df (pd.DataFrame): DataFrame with 'Price', 'Bid_Quantity', 'Ask_Quantity' columns.
        title (str): Title for the plot.
    """
    fig, ax = plt.subplots(figsize=(12, 7))

    # Filter bids and asks for clarity
    bids_df = snapshot_df[snapshot_df['Bid_Quantity'] > 0].copy()
    asks_df = snapshot_df[snapshot_df['Ask_Quantity'] > 0].copy()

    # Calculate cumulative bid volume (from highest bid downwards)
    bids_df = bids_df.sort_values(by='Price', ascending=False)
    bids_df['Cumulative_Bid_Quantity'] = bids_df['Bid_Quantity'].cumsum()

    # Calculate cumulative ask volume (from lowest ask upwards)
    asks_df = asks_df.sort_values(by='Price', ascending=True)
    asks_df['Cumulative_Ask_Quantity'] = asks_df['Ask_Quantity'].cumsum()

    # Plot cumulative bids
    ax.plot(bids_df['Price'], bids_df['Cumulative_Bid_Quantity'], drawstyle='steps-post', color='green', label='Cumulative Bids', linewidth=2)
    # Plot cumulative asks
    ax.plot(asks_df['Price'], asks_df['Cumulative_Ask_Quantity'], drawstyle='steps-pre', color='red', label='Cumulative Asks', linewidth=2)

    ax.set_title(title)
    ax.set_xlabel('Price')
    ax.set_ylabel('Cumulative Quantity')
    ax.legend()
    ax.grid(True, linestyle='--', alpha=0.6)

    # Highlight best bid and ask (if available)
    best_bid_price = bids_df['Price'].max() if not bids_df.empty else None
    best_ask_price = asks_df['Price'].min() if not asks_df.empty else None

    if best_bid_price:
        ax.axvline(best_bid_price, color='darkgreen', linestyle='--', linewidth=1, label=f'Best Bid: {best_bid_price}')
    if best_ask_price:
        ax.axvline(best_ask_price, color='darkred', linestyle='--', linewidth=1, label=f'Best Ask: {best_ask_price}')

    plt.tight_layout()
    plt.show()

plot_cumulative_order_book(full_depth_snapshot) # Use full_depth_snapshot for more comprehensive cumulative view