def top_book_imbalance(bid_volume, ask_volume):
    """Calculate imbalance ratio using best bid and best ask volumes.

    This function computes the simplest form of order book imbalance,
    considering only the volume at the best bid and best ask prices.

    Args:
        bid_volume (int): The volume of shares at the best bid price.
        ask_volume (int): The volume of shares at the best ask price.

    Returns:
        float: The imbalance ratio, ranging from -1 (strong sell) to +1 (strong buy).
               Returns 0 if total volume is zero to avoid division by zero.
    """
    if bid_volume + ask_volume == 0:
        return 0  # Avoid division by zero if there are no orders
    return (bid_volume - ask_volume) / (bid_volume + ask_volume)

def depth_imbalance(bids, asks, levels=5):
    """Calculate imbalance across multiple price levels.

    This function aggregates volumes from a specified number of top price levels
    on both the bid and ask sides to calculate a more comprehensive imbalance.

    Args:
        bids (list of tuples): A list of (price, volume) tuples for bid orders,
                               sorted by price in descending order.
        asks (list of tuples): A list of (price, volume) tuples for ask orders,
                               sorted by price in ascending order.
        levels (int): The number of top price levels to consider for aggregation.

    Returns:
        float: The imbalance ratio, ranging from -1 to +1.
               Returns 0 if aggregated total volume is zero.
    """
    # Sum the volumes for the specified number of levels
    bid_vol = sum(vol for _, vol in bids[:levels])
    ask_vol = sum(vol for _, vol in asks[:levels])

    if bid_vol + ask_vol == 0:
        return 0  # Avoid division by zero
    return (bid_vol - ask_vol) / (bid_vol + ask_vol)

def weighted_imbalance(bids, asks, levels=5):
    """Calculate imbalance where volumes are weighted by their distance from the mid price.

    This method gives more weight to orders closer to the mid-price,
    reflecting their higher probability of execution and immediate market impact.

    Args:
        bids (list of tuples): A list of (price, volume) tuples for bid orders.
        asks (list of tuples): A list of (price, volume) tuples for ask orders.
        levels (int): The number of top price levels to consider for weighting.

    Returns:
        float: The weighted imbalance ratio, ranging from -1 to +1.
               Returns 0 if there are no bids or asks, or if total weighted volume is zero.
    """
    if not bids or not asks:
        return 0 # Cannot calculate if one side is empty

    # Calculate the mid-price, which is the average of the best bid and best ask
    mid = (bids[0][0] + asks[0][0]) / 2

    # Calculate weighted bid volume: sum of (volume / distance from mid-price)
    # The division by distance gives more weight to orders closer to the mid-price.
    bid_weighted = sum(vol / abs(price - mid) for price, vol in bids[:levels] if price != mid) # Ensure price != mid to avoid ZeroDivisionError
    ask_weighted = sum(vol / abs(price - mid) for price, vol in asks[:levels] if price != mid)

    if bid_weighted + ask_weighted == 0:
        return 0  # Avoid division by zero
    return (bid_weighted - ask_weighted) / (bid_weighted + ask_weighted)

import matplotlib.pyplot as plt # Used for creating static, interactive, and animated visualizations
import numpy as np # Used for numerical operations, especially with arrays

# Define the number of price levels to visualize (e.g., top 5)
levels_to_plot = 5
# Extract volumes for the top N bid and ask levels from our mock data
bid_volumes = [vol for _, vol in bids[:levels_to_plot]]
ask_volumes = [vol for _, vol in asks[:levels_to_plot]]

# Create a figure with three subplots arranged in one row
fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(18, 5))

# --- Subplot 1 (ax1): Bar chart showing Bid vs Ask Volumes per Price Level ---
x = np.arange(levels_to_plot) # x-axis positions for the bars
width = 0.35 # Width of the bars
ax1.bar(x - width/2, bid_volumes, width, label='Bids', color='green', alpha=0.7)
ax1.bar(x + width/2, ask_volumes, width, label='Asks', color='red', alpha=0.7)
ax1.set_xlabel('Price Level (from best bid/ask)')
ax1.set_ylabel('Volume')
ax1.set_title('Bid vs Ask Volumes by Level')
ax1.set_xticks(x)
ax1.set_xticklabels([f'Level {i+1}' for i in x])
ax1.legend()
ax1.grid(axis='y', linestyle='--', alpha=0.5)

# --- Subplot 2 (ax2): Horizontal Bar chart showing Total Bid vs Total Ask Volumes ---
total_bid = sum(bid_volumes)
total_ask = sum(ask_volumes)
ax2.barh(['Total Bid', 'Total Ask'], [total_bid, total_ask], color=['green', 'red'], alpha=0.7)
ax2.set_xlabel('Total Volume')
ax2.set_title(f'Total Volume Comparison (Depth Imbalance: {depth_imb:.3f})')
for i, v in enumerate([total_bid, total_ask]):
    ax2.text(v + 20, i, str(v), va='center', ha='left', fontsize=10)
ax2.set_xlim(0, max(total_bid, total_ask) * 1.2)

# --- Subplot 3 (ax3): Gauge Meter for Imbalance Direction ---
imbalance_value = depth_imb # Use the depth imbalance calculated earlier
norm_imb = (imbalance_value + 1) / 2 # Normalize the imbalance value from [-1, 1] to [0, 1] for plotting on a bar

ax3.barh(['Imbalance'], [1], color='lightgray', alpha=0.5, height=0.5)
ax3.barh(['Imbalance'], [norm_imb], color='green' if imbalance_value > 0 else 'red', alpha=0.8, height=0.5)
ax3.set_xlim(0, 1)
ax3.set_xticks([0, 0.5, 1])
ax3.set_xticklabels(['Strong Sell (-1)', 'Balanced (0)', 'Strong Buy (+1)'])
ax3.set_title(f'Imbalance Gauge: {imbalance_value:.3f}')
ax3.axvline(x=0.5, color='black', linestyle='--', linewidth=0.8)
ax3.get_yaxis().set_visible(False)

plt.tight_layout() # Adjust subplot params for a tight layout for all three subplots
plt.show() # Display the first set of plots

# Calculate cumulative bid and ask volumes for each level
cum_bid = np.cumsum(bid_volumes)
cum_ask = np.cumsum(ask_volumes)

# Calculate cumulative imbalance at each level
cum_imb = (cum_bid - cum_ask) / (cum_bid + cum_ask)

plt.figure(figsize=(10, 5))
plt.plot(range(1, levels_to_plot + 1), cum_imb, marker='o', linewidth=2, color='purple')
plt.axhline(y=0, color='black', linestyle='--', linewidth=0.7, label='Balanced (0)')
plt.axhline(y=0.3, color='green', linestyle=':', alpha=0.7, label='Buy Pressure Threshold (+0.3)')
plt.axhline(y=-0.3, color='red', linestyle=':', alpha=0.7, label='Sell Pressure Threshold (-0.3)')
plt.xlabel('Number of Levels Included (Depth)')
plt.ylabel('Cumulative Imbalance Ratio')
plt.title('Cumulative Imbalance as More Order Book Depth is Considered')
plt.grid(True, alpha=0.4)
plt.xticks(range(1, levels_to_plot + 1))
plt.legend()
plt.show() # Display the second plot

# Mock order book data for demonstration
# Bids: (price, volume) - sorted from highest price to lowest
# Asks: (price, volume) - sorted from lowest price to highest
bids = [(100.00, 500), (99.99, 300), (99.98, 200), (99.97, 150), (99.96, 100)]
asks = [(100.01, 100), (100.02, 150), (100.03, 200), (100.04, 250), (100.05, 300)]

# Extracting best bid and ask volumes for the Top-of-Book Imbalance calculation
best_bid_vol = bids[0][1]   # Volume at the highest bid price (100.00, 500 shares)
best_ask_vol = asks[0][1]   # Volume at the lowest ask price (100.01, 100 shares)

# Calculate all three types of imbalances using the defined functions
top_imb = top_book_imbalance(best_bid_vol, best_ask_vol)
depth_imb = depth_imbalance(bids, asks, levels=5) # Using top 5 levels for depth imbalance
weighted_imb = weighted_imbalance(bids, asks, levels=5) # Using top 5 levels for weighted imbalance

# Print the calculated imbalance values, formatted to 3 decimal places
print(f"Top-of-book imbalance: {top_imb:.3f}")
print(f"Depth imbalance (5 levels): {depth_imb:.3f}")
print(f"Weighted imbalance: {weighted_imb:.3f}")

# Interpret the Top-of-Book Imbalance result based on the thresholds introduced earlier
print("\nInterpretation (based on Top-of-Book Imbalance):")
if top_imb > 0.3:
    print("Strong buy imbalance → price likely to go up")
elif top_imb < -0.3:
    print("Strong sell imbalance → price likely to go down")
else:
    print("Balanced book → buying and selling pressures are relatively even")