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

def generate_mock_trade_data(num_data_points=100, base_price=100.0, seed=42):
    """
    Generates mock high-frequency trading data.

    Args:
        num_data_points (int): Number of data points to generate.
        base_price (float): Base price for the simulated asset.
        seed (int): Random seed for reproducibility.

    Returns:
        pd.DataFrame: A DataFrame containing simulated trade data.
    """
    np.random.seed(seed)

    time_stamps = pd.to_datetime(pd.date_range(start='2023-01-01 09:00:00', periods=num_data_points, freq='1S'))

    price_changes = np.random.normal(0, 0.05, num_data_points).cumsum()
    price = base_price + price_changes

    bid_price = price - np.random.uniform(0.005, 0.01, num_data_points)
    ask_price = price + np.random.uniform(0.005, 0.01, num_data_points)

    bid_volume = np.random.randint(50, 500, num_data_points)
    ask_volume = np.random.randint(50, 500, num_data_points)

    aggressive_buy_volume = np.random.randint(0, 200, num_data_points)
    aggressive_sell_volume = np.random.randint(0, 200, num_data_points)

    aggressive_buy_volume[30:70] = aggressive_buy_volume[30:70] + np.random.randint(50, 150, 40)
    aggressive_sell_volume[70:90] = aggressive_sell_volume[70:90] + np.random.randint(50, 100, 20)

    df = pd.DataFrame({
        'timestamp': time_stamps,
        'price': price,
        'bid_price': bid_price,
        'ask_price': ask_price,
        'bid_volume': bid_volume,
        'ask_volume': ask_volume,
        'aggressive_buy_volume': aggressive_buy_volume,
        'aggressive_sell_volume': aggressive_sell_volume
    })

    df = df.set_index('timestamp')
    return df

# --- 1. Mock Data Generation ---
df = generate_mock_trade_data()

print("Mock Data Head:")
display(df.head())

def calculate_bid_ask_volume_imbalance(df: pd.DataFrame) -> pd.Series:
    """
    Calculates the Bid/Ask Volume Imbalance.

    Formula: Total Bid Volume - Total Ask Volume

    Args:
        df (pd.DataFrame): DataFrame with 'bid_volume' and 'ask_volume' columns.

    Returns:
        pd.Series: A Series containing the Bid/Ask Volume Imbalance.
    """
    if 'bid_volume' not in df.columns or 'ask_volume' not in df.columns:
        raise ValueError("DataFrame must contain 'bid_volume' and 'ask_volume' columns.")

    return df['bid_volume'] - df['ask_volume']


def calculate_order_flow_imbalance(df: pd.DataFrame) -> pd.Series:
    """
    Calculates the Order Flow Imbalance (OFI).

    Formula: Aggressive Buy Volume - Aggressive Sell Volume

    Args:
        df (pd.DataFrame): DataFrame with 'aggressive_buy_volume' and 'aggressive_sell_volume' columns.

    Returns:
        pd.Series: A Series containing the Order Flow Imbalance.
    """
    if 'aggressive_buy_volume' not in df.columns or 'aggressive_sell_volume' not in df.columns:
        raise ValueError("DataFrame must contain 'aggressive_buy_volume' and 'aggressive_sell_volume' columns.")

    return df['aggressive_buy_volume'] - df['aggressive_sell_volume']


def calculate_cumulative_volume_delta(order_flow_imbalance: pd.Series) -> pd.Series:
    """
    Calculates the Cumulative Volume Delta (CVD).

    Formula: CVD_t = CVD_{t-1} + OFI_t

    Args:
        order_flow_imbalance (pd.Series): A Series containing the Order Flow Imbalance values.

    Returns:
        pd.Series: A Series containing the Cumulative Volume Delta.
    """
    return order_flow_imbalance.cumsum()


# --- 2. Perform Calculations with Mock Data ---

# Bid/Ask Volume Imbalance
df['bid_ask_volume_imbalance'] = calculate_bid_ask_volume_imbalance(df)
print("\nBid/Ask Volume Imbalance (first 5 values):")
display(df['bid_ask_volume_imbalance'].head())
print(f"Interpretation: A positive value like {df['bid_ask_volume_imbalance'].iloc[0]:.2f} indicates more bid liquidity than ask liquidity at that moment. A negative value suggests the opposite.")

# Order Flow Imbalance (OFI)
df['order_flow_imbalance'] = calculate_order_flow_imbalance(df)
print("\nOrder Flow Imbalance (first 5 values):")
display(df['order_flow_imbalance'].head())
print(f"Interpretation: A positive value like {df['order_flow_imbalance'].iloc[0]:.2f} means there was more aggressive buying than selling in that second. A negative value means more aggressive selling.")

# Cumulative Volume Delta (CVD)
df['cumulative_volume_delta'] = calculate_cumulative_volume_delta(df['order_flow_imbalance'])
print("\nCumulative Volume Delta (first 5 values):")
display(df['cumulative_volume_delta'].head())
print(f"Interpretation: A cumulative value like {df['cumulative_volume_delta'].iloc[4]:.2f} shows the net accumulated aggressive buying/selling pressure up to that point. It's a trend indicator.")

import matplotlib.ticker as mticker

# --- Visualization 1: Multi-panel plot of Imbalances and Price ---

fig, axes = plt.subplots(3, 1, figsize=(14, 12), sharex=True)
fig.suptitle('Trade Flow Imbalance Analysis', fontsize=16)

# Panel 1: Price
axes[0].plot(df.index, df['price'], label='Price', color='blue', linewidth=1.5)
axes[0].set_title('Price Trend')
axes[0].set_ylabel('Price')
axes[0].grid(True, linestyle='--', alpha=0.7)
axes[0].legend()
# Adjust y-axis limits for price to zoom in
price_min = df['price'].min()
price_max = df['price'].max()
axes[0].set_ylim(price_min - (price_max - price_min) * 0.1, price_max + (price_max - price_min) * 0.1)

# Panel 2: Bid/Ask Volume Imbalance
axes[1].plot(df.index, df['bid_ask_volume_imbalance'], label='Bid/Ask Volume Imbalance', color='green', alpha=0.8)
axes[1].axhline(0, color='red', linestyle='--', linewidth=0.8)
axes[1].set_title('Bid/Ask Volume Imbalance (Liquidity Pressure)')
axes[1].set_ylabel('Imbalance')
axes[1].grid(True, linestyle='--', alpha=0.7)
axes[1].legend()
# Adjust y-axis limits for bid/ask volume imbalance to zoom in
imbalance_max_abs = df['bid_ask_volume_imbalance'].abs().max()
axes[1].set_ylim(-imbalance_max_abs * 1.1, imbalance_max_abs * 1.1)

# Panel 3: Order Flow Imbalance
axes[2].bar(df.index, df['order_flow_imbalance'], label='Order Flow Imbalance', color=np.where(df['order_flow_imbalance'] > 0, 'purple', 'orange'), width=1.0)
axes[2].axhline(0, color='red', linestyle='--', linewidth=0.8)
axes[2].set_title('Order Flow Imbalance (Aggressive Pressure)')
axes[2].set_ylabel('Imbalance')
axes[2].set_xlabel('Time')
axes[2].grid(True, linestyle='--', alpha=0.7)
axes[2].legend()

# Explicitly set x-axis limits to match data range
axes[0].set_xlim(df.index.min(), df.index.max())

# Format x-axis for better readability
plt.tight_layout(rect=[0, 0.03, 1, 0.96]) # Adjust layout to prevent title overlap
plt.show()

# --- Visualization 2: Cumulative Volume Delta (CVD) vs. Price ---

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

ax1.set_xlabel('Time')
ax1.set_ylabel('Price', color='blue')
ax1.plot(df.index, df['price'], color='blue', label='Price', linewidth=2)
ax1.tick_params(axis='y', labelcolor='blue')
ax1.grid(True, linestyle='--', alpha=0.7)

# Create a second y-axis for CVD
ax2 = ax1.twinx()
ax2.set_ylabel('Cumulative Volume Delta', color='red')
ax2.plot(df.index, df['cumulative_volume_delta'], color='red', linestyle='--', label='Cumulative Volume Delta', alpha=0.8)
ax2.tick_params(axis='y', labelcolor='red')

# Add a threshold for CVD, e.g., 2000 units
cvd_threshold_positive = 2000
cvd_threshold_negative = -2000
ax2.axhline(cvd_threshold_positive, color='gray', linestyle=':', label=f'CVD Threshold (+{cvd_threshold_positive})')
ax2.axhline(cvd_threshold_negative, color='gray', linestyle=':', label=f'CVD Threshold ({cvd_threshold_negative})')

fig.suptitle('Price vs. Cumulative Volume Delta (CVD)', fontsize=16)
fig.legend(loc="upper left", bbox_to_anchor=(0.1, 0.9))
plt.tight_layout(rect=[0, 0.03, 1, 0.96]) # Adjust layout
plt.show()