import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.figure import Figure
from matplotlib.axes import Axes
from typing import Tuple, Optional
import warnings

warnings.filterwarnings('ignore')

# ── Matplotlib Style Configuration ──────────────────────────────────────────
plt.rcParams.update({
    'figure.facecolor':  '#FFFFFF',
    'axes.facecolor':    '#FFFFFF',
    'axes.edgecolor':    '#30363d',
    'axes.labelcolor':   '#c9d1d9',
    'axes.grid':         True,
    'grid.color':        '#21262d',
    'grid.linewidth':    0.8,
    'xtick.color':       '#8b949e',
    'ytick.color':       '#8b949e',
    'text.color':        '#c9d1d9',
    'legend.facecolor':  '#161b22',
    'legend.edgecolor':  '#30363d',
    'font.family':       'monospace',
    'font.size':         10,
})

print("Imports and style configuration complete.")

def generate_sample_order_book(
    mid_price: float = 100.0,
    num_levels: int = 20,
    tick_size: float = 0.10,
    base_volume: float = 1000.0,
    seed: Optional[int] = 42,
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """
    Generate a synthetic order book with realistic bid/ask price levels.

    The volume distribution uses exponential decay from the midpoint outward,
    with multiplicative noise to simulate real-world irregularity. Occasional
    'liquidity walls' (large volume spikes) are injected to mimic iceberg-style
    resting orders often found at round-number price levels.

    Parameters
    ----------
    mid_price   : Reference midpoint price.
    num_levels  : Number of price levels per side.
    tick_size   : Price increment between consecutive levels.
    base_volume : Approximate maximum volume at the innermost level.
    seed        : Random seed for reproducibility.

    Returns
    -------
    bids_df : DataFrame[price, volume] — sorted descending (best bid first).
    asks_df : DataFrame[price, volume] — sorted ascending (best ask first).

    Raises
    ------
    ValueError
        If num_levels < 1, tick_size <= 0, or base_volume <= 0.
    """
    if num_levels < 1:
        raise ValueError(f"num_levels must be >= 1, got {num_levels}")
    if tick_size <= 0:
        raise ValueError(f"tick_size must be > 0, got {tick_size}")
    if base_volume <= 0:
        raise ValueError(f"base_volume must be > 0, got {base_volume}")

    rng = np.random.default_rng(seed)

    half_spread = tick_size / 2.0

    # ── Price levels ─────────────────────────────────────────────────────────
    bid_prices = np.array([
        round(mid_price - half_spread - i * tick_size, 10)
        for i in range(num_levels)
    ])
    ask_prices = np.array([
        round(mid_price + half_spread + i * tick_size, 10)
        for i in range(num_levels)
    ])

    # ── Volume distribution: exponential decay + multiplicative noise
    decay = np.exp(-0.15 * np.arange(num_levels))
    noise_bids = rng.lognormal(mean=0.0, sigma=0.35, size=num_levels)
    noise_asks = rng.lognormal(mean=0.0, sigma=0.35, size=num_levels)

    bid_volumes = np.maximum(1.0, base_volume * decay * noise_bids)
    ask_volumes = np.maximum(1.0, base_volume * decay * noise_asks)

    # ── Inject liquidity walls at 2–3 random levels per side
    wall_indices_bid = rng.choice(num_levels, size=min(3, num_levels), replace=False)
    wall_indices_ask = rng.choice(num_levels, size=min(3, num_levels), replace=False)
    bid_volumes[wall_indices_bid] *= rng.uniform(2.5, 5.0, size=len(wall_indices_bid))
    ask_volumes[wall_indices_ask] *= rng.uniform(2.5, 5.0, size=len(wall_indices_ask))

    # ── Build DataFrames
    bids_df = pd.DataFrame({'price': bid_prices, 'volume': bid_volumes})
    asks_df = pd.DataFrame({'price': ask_prices, 'volume': ask_volumes})

    return bids_df, asks_df


# ── Quick sanity check
_bids_test, _asks_test = generate_sample_order_book()
print(f"{'Bids':─^40}")
print(_bids_test.head(5).to_string(index=False))
print(f"\n{'Asks':─^40}")
print(_asks_test.head(5).to_string(index=False))
print(f"\ngenerate_sample_order_book() — {len(_bids_test)} bid levels, {len(_asks_test)} ask levels")

def compute_cumulative_depth(
    bids_df: pd.DataFrame,
    asks_df: pd.DataFrame,
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """
    Compute cumulative volume depth for bid and ask sides of the order book.

    Bids are sorted descending (best bid first) and cumulated from index 0
    downward in price. Asks are sorted ascending (best ask first) and cumulated
    from index 0 upward in price. The result is the classic staircase shape
    seen in exchange depth charts.

    Parameters
    ----------
    bids_df : DataFrame with 'price' and 'volume' columns (bid side).
    asks_df : DataFrame with 'price' and 'volume' columns (ask side).

    Returns
    -------
    bids_cum : Copy of bids_df with added 'cumulative_volume' column.
    asks_cum : Copy of asks_df with added 'cumulative_volume' column.

    Raises
    ------
    ValueError
        If either DataFrame is missing required columns or is empty.
    """
    required_cols = {'price', 'volume'}

    for name, df in [('bids_df', bids_df), ('asks_df', asks_df)]:
        if df.empty:
            raise ValueError(f"{name} is empty.")
        missing = required_cols - set(df.columns)
        if missing:
            raise ValueError(f"{name} missing columns: {missing}")
        if df['volume'].lt(0).any():
            raise ValueError(f"{name} contains negative volume values.")

    # ── Bids: descending price order, cumulate from best bid downward
    bids_cum = (
        bids_df
        .copy()
        .sort_values('price', ascending=False)
        .reset_index(drop=True)
    )
    bids_cum['cumulative_volume'] = bids_cum['volume'].cumsum()

    # ── Asks: ascending price order, cumulate from best ask upward
    asks_cum = (
        asks_df
        .copy()
        .sort_values('price', ascending=True)
        .reset_index(drop=True)
    )
    asks_cum['cumulative_volume'] = asks_cum['volume'].cumsum()

    return bids_cum, asks_cum


# ── Quick sanity check
_bids_cum, _asks_cum = compute_cumulative_depth(_bids_test, _asks_test)
print(f"{'Bids with Cumulative Depth':─^55}")
print(_bids_cum[['price', 'volume', 'cumulative_volume']].head(5).to_string(index=False))
print(f"\n{'Asks with Cumulative Depth':─^55}")
print(_asks_cum[['price', 'volume', 'cumulative_volume']].head(5).to_string(index=False))
print(f"\ncompute_cumulative_depth() — max bid depth: {_bids_cum['cumulative_volume'].max():.1f}, "
      f"max ask depth: {_asks_cum['cumulative_volume'].max():.1f}")

def prepare_depth_data(
    bids_cum: pd.DataFrame,
    asks_cum: pd.DataFrame,
) -> dict:
    """
    Merge cumulative bid/ask data and compute market microstructure metrics.

    Derives best bid, best ask, midpoint, spread (absolute and relative),
    and the bid/ask volume imbalance ratio — all essential metrics for
    interpreting order book state and annotating the depth chart.

    Parameters
    ----------
    bids_cum : Bids DataFrame containing 'price' and 'cumulative_volume'.
    asks_cum : Asks DataFrame containing 'price' and 'cumulative_volume'.

    Returns
    -------
    dict with keys:
        bids            : pd.DataFrame
        asks            : pd.DataFrame
        best_bid        : float
        best_ask        : float
        mid_price       : float
        spread          : float
        spread_pct      : float   (spread / mid_price * 100)
        bid_ask_imbalance : float  (range [-1, +1]; positive = more bid-side volume)

    Raises
    ------
    ValueError
        If DataFrames are missing required columns or best bid >= best ask
        (crossed book, indicating upstream data error).
    """
    for name, df in [('bids_cum', bids_cum), ('asks_cum', asks_cum)]:
        required = {'price', 'cumulative_volume'}
        missing = required - set(df.columns)
        if missing:
            raise ValueError(f"{name} missing columns: {missing}")
        if df.empty:
            raise ValueError(f"{name} is empty.")

    best_bid: float = bids_cum['price'].max()
    best_ask: float = asks_cum['price'].min()

    if best_bid >= best_ask:
        raise ValueError(
            f"Crossed book detected: best_bid ({best_bid}) >= best_ask ({best_ask}). "
            "Verify upstream data generation."
        )

    mid_price: float = (best_bid + best_ask) / 2.0
    spread: float = best_ask - best_bid
    spread_pct: float = (spread / mid_price) * 100.0

    # ── Bid-Ask Imbalance: (BidVol - AskVol) / (BidVol + AskVol)
    total_bid_vol: float = bids_cum['cumulative_volume'].max()
    total_ask_vol: float = asks_cum['cumulative_volume'].max()
    denom = total_bid_vol + total_ask_vol
    bid_ask_imbalance: float = (
        (total_bid_vol - total_ask_vol) / denom if denom > 0 else 0.0
    )

    return {
        'bids':               bids_cum,
        'asks':               asks_cum,
        'best_bid':           best_bid,
        'best_ask':           best_ask,
        'mid_price':          mid_price,
        'spread':             spread,
        'spread_pct':         spread_pct,
        'bid_ask_imbalance':  bid_ask_imbalance,
    }


# ── Quick sanity check
_plot_data = prepare_depth_data(_bids_cum, _asks_cum)
print(f"{'Market Microstructure Metrics':─^55}")
print(f"  Best Bid     : {_plot_data['best_bid']:>10.4f}")
print(f"  Best Ask     : {_plot_data['best_ask']:>10.4f}")
print(f"  Mid Price    : {_plot_data['mid_price']:>10.4f}")
print(f"  Spread       : {_plot_data['spread']:>10.4f}")
print(f"  Spread (%)   : {_plot_data['spread_pct']:>10.4f}%")
print(f"  Imbalance    : {_plot_data['bid_ask_imbalance']:>10.4f}  (>0 = bid-heavy)")
print(f"\nprepare_depth_data() — all metrics computed successfully")

def plot_order_book_depth(
    plot_data: dict,
    title: str = 'Order Book Depth',
    figsize: Tuple[float, float] = (14, 7),
    show_volume_bars: bool = True,
) -> Tuple[Figure, Axes]:
    """
    Render a professional, annotated order book depth chart.

    Plots cumulative bid and ask volumes as step curves, fills below each
    curve with transparent colour, highlights the bid-ask spread region,
    marks the midpoint price, and overlays a microstructure metrics box.
    Optionally displays per-level volume bars for granular liquidity insight.

    Parameters
    ----------
    plot_data        : dict returned by prepare_depth_data().
    title            : Chart title string.
    figsize          : (width, height) in inches.
    show_volume_bars : If True, draw translucent volume bars at each price level.

    Returns
    -------
    fig : matplotlib Figure object.
    ax  : matplotlib Axes object.

    Raises
    ------
    KeyError
        If plot_data is missing required keys.
    """
    required_keys = {'bids', 'asks', 'best_bid', 'best_ask', 'mid_price',
                     'spread', 'spread_pct', 'bid_ask_imbalance'}
    missing_keys = required_keys - set(plot_data.keys())
    if missing_keys:
        raise KeyError(f"plot_data missing keys: {missing_keys}")

    bids:     pd.DataFrame = plot_data['bids']
    asks:     pd.DataFrame = plot_data['asks']
    best_bid: float        = plot_data['best_bid']
    best_ask: float        = plot_data['best_ask']
    mid_price: float       = plot_data['mid_price']
    spread:    float       = plot_data['spread']
    spread_pct: float      = plot_data['spread_pct']
    imbalance: float       = plot_data['bid_ask_imbalance']

    # ── Colour palette
    BID_COLOR    = '#00d97e'   # green
    ASK_COLOR    = '#ff4d6d'   # red
    SPREAD_COLOR = '#f7c948'   # amber
    MID_COLOR    = '#a5b4fc'   # lavender
    BG_COLOR     = '#0d1117'

    fig, ax = plt.subplots(figsize=figsize, facecolor=BG_COLOR)
    ax.set_facecolor(BG_COLOR)

    # ── Bid depth curve
    ax.step(
        bids['price'], bids['cumulative_volume'],
        where='post', color=BID_COLOR, linewidth=2.0,
        label='Cumulative Bid Depth', zorder=4,
    )
    ax.fill_between(
        bids['price'], bids['cumulative_volume'],
        step='post', alpha=0.18, color=BID_COLOR, zorder=3,
    )

    # ── Ask depth curve
    ax.step(
        asks['price'], asks['cumulative_volume'],
        where='pre', color=ASK_COLOR, linewidth=2.0,
        label='Cumulative Ask Depth', zorder=4,
    )
    ax.fill_between(
        asks['price'], asks['cumulative_volume'],
        step='pre', alpha=0.18, color=ASK_COLOR, zorder=3,
    )

    # ── Optional per-level volume bars
    if show_volume_bars:
        tick_size_est = abs(bids['price'].diff().median())
        bar_width = tick_size_est * 0.6 if tick_size_est > 0 else 0.05

        ax.bar(
            bids['price'], bids['volume'],
            width=bar_width, color=BID_COLOR, alpha=0.35,
            align='center', zorder=2, label='Bid Volume per Level',
        )
        ax.bar(
            asks['price'], asks['volume'],
            width=bar_width, color=ASK_COLOR, alpha=0.35,
            align='center', zorder=2, label='Ask Volume per Level',
        )

    # ── Spread region shading
    y_max = max(bids['cumulative_volume'].max(), asks['cumulative_volume'].max())
    ax.axvspan(
        best_bid, best_ask,
        alpha=0.12, color=SPREAD_COLOR, zorder=1,
        label=f'Spread ({spread_pct:.3f}%)',
    )

    # ── Midpoint price marker
    ax.axvline(
        mid_price, color=MID_COLOR, linewidth=1.5,
        linestyle='--', alpha=0.85, zorder=5,
        label=f'Mid Price ({mid_price:.2f})',
    )
    ax.text(
        mid_price, y_max * 0.97,
        f' Mid\n {mid_price:.2f}',
        color=MID_COLOR, fontsize=8.5, va='top',
        fontweight='bold', zorder=6,
    )

    # ── Best bid/ask tick markers
    # Calculate an estimated tick size for positioning labels
    tick_size_est = abs(bids['price'].diff().median())

    label_offset = tick_size_est * 1 # A couple of tick sizes away

    for price, color, label, h_align in [
        (best_bid, BID_COLOR, f'Best Bid\n{best_bid:.2f}', 'right'),
        (best_ask, ASK_COLOR, f'Best Ask\n{best_ask:.2f}', 'left'),
    ]:
        # Adjust price for text placement
        text_price = price - label_offset if h_align == 'right' else price + label_offset

        ax.axvline(price, color=color, linewidth=1.0,
                   linestyle=':', alpha=0.6, zorder=5)
        ax.text(
            text_price, y_max * 0.85, label,
            color=color, fontsize=7.5, ha=h_align, # Use h_align here
            va='top', alpha=0.9, zorder=6,
        )

    # ── Microstructure metrics annotation box
    imb_sign = '+' if imbalance > 0 else ''
    imb_color = BID_COLOR if imbalance > 0 else ASK_COLOR if imbalance < 0 else '#8b949e'
    stats_text = (
        f"Spread:     {spread:.4f}  ({spread_pct:.4f}%)\n"
        f"Best Bid:   {best_bid:.4f}\n"
        f"Best Ask:   {best_ask:.4f}\n"
        f"Mid Price:  {mid_price:.4f}\n"
        f"Imbalance:  {imb_sign}{imbalance:.4f}"
    )
    ax.text(
        0.01, 0.98, stats_text,
        transform=ax.transAxes, fontsize=8,
        verticalalignment='top', fontfamily='monospace',
        bbox=dict(boxstyle='round,pad=0.5', facecolor='#161b22',
                  edgecolor='#30363d', alpha=0.90),
        color='#c9d1d9', zorder=10,
    )

    # ── Axes formatting
    ax.set_xlabel('Price', fontsize=12, labelpad=10, color='#8b949e')
    ax.set_ylabel('Cumulative Volume', fontsize=12, labelpad=10, color='#8b949e')
    ax.set_title(title, fontsize=15, fontweight='bold', pad=18,
                 color='#f0f6fc', fontfamily='monospace')

    ax.set_xlim(
        bids['price'].min() - spread,
        asks['price'].max() + spread,
    )
    ax.set_ylim(bottom=0, top=y_max * 1.10)

    ax.tick_params(axis='both', which='major', labelsize=8.5,
                   colors='#8b949e')
    ax.spines[['top', 'right']].set_visible(False)
    ax.spines[['left', 'bottom']].set_color('#30363d')

    ax.yaxis.set_major_formatter(
        plt.FuncFormatter(lambda v, _: f'{v:,.0f}')
    )

    legend = ax.legend(
        loc='upper right', fontsize=8.5,
        framealpha=0.85, ncol=2,
    )
    for text in legend.get_texts():
        text.set_color('#c9d1d9')

    # ── Watermark
    fig.text(
        0.99, 0.01,
        'Market Microstructure Series — Order Book Depth',
        ha='right', va='bottom', fontsize=7,
        color='#30363d', style='italic',
    )

    plt.tight_layout()
    return fig, ax

def main(
    mid_price: float = 100.0,
    num_levels: int = 25,
    tick_size: float = 0.10,
    base_volume: float = 1500.0,
    seed: Optional[int] = 42,
    chart_title: str = 'Order Book Depth — Synthetic Exchange',
) -> Tuple[Figure, Axes, dict]:
    """
    Execute the complete order book depth visualisation pipeline.

    Pipeline stages
    ---------------
    1. generate_sample_order_book()  — create synthetic bid/ask levels
    2. compute_cumulative_depth()    — compute running volume totals
    3. prepare_depth_data()          — merge & compute microstructure metrics
    4. plot_order_book_depth()       — render annotated depth chart

    Parameters
    ----------
    mid_price    : Reference midpoint price for order book generation.
    num_levels   : Number of price levels on each side.
    tick_size    : Minimum price increment between levels.
    base_volume  : Approximate peak volume at the innermost level.
    seed         : Random seed for reproducibility.
    chart_title  : Title displayed on the depth chart.

    Returns
    -------
    fig       : matplotlib Figure.
    ax        : matplotlib Axes.
    plot_data : dict of prepared depth data and microstructure metrics.
    """
    print('─' * 60)
    print('  Order Book Depth Visualiser')
    print('  Market Microstructure Series')
    print('─' * 60)

    # ── Stage 1: Generate synthetic order book
    print('\n[1/4] Generating synthetic order book...')
    bids_raw, asks_raw = generate_sample_order_book(
        mid_price=mid_price,
        num_levels=num_levels,
        tick_size=tick_size,
        base_volume=base_volume,
        seed=seed,
    )
    print(f'     Bid levels: {len(bids_raw):>4}  |  '
          f'Ask levels: {len(asks_raw):>4}  |  '
          f'Price range: [{bids_raw["price"].min():.3f}, {asks_raw["price"].max():.3f}]')

    # ── Stage 2: Compute cumulative depth
    print('\n[2/4] Computing cumulative depth...')
    bids_cum, asks_cum = compute_cumulative_depth(bids_raw, asks_raw)
    print(f'     Max bid depth: {bids_cum["cumulative_volume"].max():>10,.1f}')
    print(f'     Max ask depth: {asks_cum["cumulative_volume"].max():>10,.1f}')

    # ── Stage 3: Prepare plot data
    print('\n[3/4] Preparing plot data and computing metrics...')
    plot_data = prepare_depth_data(bids_cum, asks_cum)
    print(f'     Mid Price  : {plot_data["mid_price"]:>10.4f}')
    print(f'     Spread     : {plot_data["spread"]:>10.4f}  ({plot_data["spread_pct"]:.4f}%)')
    print(f'     Imbalance  : {plot_data["bid_ask_imbalance"]:>+10.4f}  '
          f'({"bid-heavy" if plot_data["bid_ask_imbalance"] > 0 else "ask-heavy" if plot_data["bid_ask_imbalance"] < 0 else "balanced"})')

    # ── Stage 4: Render depth chart
    print('\n[4/4] Rendering depth chart...')
    fig, ax = plot_order_book_depth(
        plot_data,
        title=chart_title,
        figsize=(14, 7),
        show_volume_bars=True,
    )
    plt.show()

    print('\nPipeline complete.')
    print('─' * 60)

    return fig, ax, plot_data


print("main() defined — execute the next cell to run the full pipeline.")

fig, ax, data = main(
    mid_price   = 100.0,
    num_levels  = 25,
    tick_size   = 0.10,
    base_volume = 1500.0,
    seed        = 42,
    chart_title = 'Order Book Depth — Synthetic Exchange',
)

fig_btc, ax_btc, data_btc = main(
    mid_price   = 67_500.0,
    num_levels  = 30,
    tick_size   = 10.0,
    base_volume = 5.0,        # BTC units (realistic for top-tier exchange)
    seed        = 7,
    chart_title = 'Order Book Depth — BTC/USDT (Simulated)',
)