"""Order Book Reconstruction from NASDAQ ITCH Messages — build limit order book from MBO data."""

import os
import warnings
from datetime import datetime

warnings.filterwarnings("ignore")

import matplotlib.pyplot as plt
import numpy as np
import polars as pl
from limit_orderbook import (
    get_stock_locate_mapping,
    load_itch_messages,
    reconstruct_lob_with_ofi,
)

from data.equities.loader import load_nasdaq_itch
from utils.paths import get_output_dir

# Input: Pre-parsed ITCH messages from canonical data location
ITCH_DIR = load_nasdaq_itch(get_base_path=True)

# Output: LOB snapshots organized by symbol
OUTPUT_DIR = get_output_dir(3, "nasdaq_itch") / "order_book"
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

print(f"Input messages: {ITCH_DIR}")
print(f"Output directory: {OUTPUT_DIR}")

# Configuration — Papermill injects overrides for CI
SYMBOL = "AAPL"
TRADING_DATE = "2020-01-30"
START_TIME = "09:30:00"
END_TIME = "16:00:00"
MAX_MESSAGES = 0  # 0 = all messages

# Symbol can be overridden via environment variable for batch processing
SYMBOL = os.environ.get("ITCH_SYMBOL", SYMBOL)

# Normalize MAX_MESSAGES: 0 means no limit
if MAX_MESSAGES == 0:
    MAX_MESSAGES = None

# Validate parsed ITCH data — produced by 01_itch_parser or Rust parser
assert ITCH_DIR.exists(), (
    f"Parsed ITCH data not found at {ITCH_DIR}.\n"
    "Run the ITCH pipeline first:\n"
    "  1. Download: uv run python data/equities/market/microstructure/nasdaq_itch_download.py\n"
    "  2. Parse:    Run 01_itch_parser.py (Section 4) or Rust parser (Section 6)"
)
msg_types = sorted([d.name for d in ITCH_DIR.iterdir() if d.is_dir()])
assert len(msg_types) > 0, f"No message types in {ITCH_DIR} — run 01_itch_parser first"
print(f"Available message types: {msg_types}")

# Get the stock_locate ID for our symbol from R messages
stock_map = get_stock_locate_mapping(ITCH_DIR)
assert SYMBOL in stock_map, f"Symbol {SYMBOL} not found in stock directory"
symbol_locate = stock_map[SYMBOL]
print(f"Symbol {SYMBOL} has stock_locate = {symbol_locate}")

# Add orders (A and F types) - have 'stock' column
add_a = load_itch_messages(ITCH_DIR, "A", symbol=SYMBOL, max_messages=MAX_MESSAGES)
add_f = load_itch_messages(ITCH_DIR, "F", symbol=SYMBOL, max_messages=MAX_MESSAGES)

# Combine A and F (select common columns)
common_cols = [
    "stock_locate",
    "tracking_number",
    "timestamp",
    "order_reference_number",
    "buy_sell_indicator",
    "shares",
    "stock",
    "price",
]
add_orders = pl.concat([add_a.select(common_cols), add_f.select(common_cols)])

# D, X, E, C, U messages don't have 'stock' column - use stock_locate filtering
deletes = load_itch_messages(ITCH_DIR, "D", stock_locate=symbol_locate, max_messages=MAX_MESSAGES)
cancels = load_itch_messages(ITCH_DIR, "X", stock_locate=symbol_locate, max_messages=MAX_MESSAGES)
executions = load_itch_messages(
    ITCH_DIR, "E", stock_locate=symbol_locate, max_messages=MAX_MESSAGES
)
executions_c = load_itch_messages(
    ITCH_DIR, "C", stock_locate=symbol_locate, max_messages=MAX_MESSAGES
)
replaces = load_itch_messages(ITCH_DIR, "U", stock_locate=symbol_locate, max_messages=MAX_MESSAGES)

# P messages (trades) have 'stock' column
trades = load_itch_messages(ITCH_DIR, "P", symbol=SYMBOL, max_messages=MAX_MESSAGES)

# =========================================================================
# CRITICAL: Pre-join D/X/E messages with Add orders to get price and side
# =========================================================================
# This is essential for correct LOB reconstruction. D/X/E messages only have
# order_reference_number - they need to be joined with Add orders to get
# the price and side for each order.

order_cols = ["order_reference_number", "buy_sell_indicator", "price", "shares"]
orders_lookup = add_orders.select(order_cols).rename({"shares": "original_shares"})

# =========================================================================
# IMPORTANT: Do NOT filter D/X/E/C messages based on order lookup!
# =========================================================================
# Orders can be created by U (Replace) messages, not just A/F. The chain:
#   A → U → U → E → D
# means E and D reference an order created by U, which isn't in orders_lookup.
# We look up order info from submitted_orders during reconstruction anyway.
# The join here is ONLY for backward compatibility with code that expects
# these columns - but we DON'T filter based on join results.
#
# Note: The original pre-join approach assumed all orders start with A/F,
# which is wrong for ~95% of replace chains.

# Join D messages (but DON'T filter - order may be from U chain)
if len(deletes) > 0:
    deletes = deletes.join(
        orders_lookup,
        on="order_reference_number",
        how="left",
    )
    # NOT filtered - order info looked up from pool during reconstruction

# Join X messages (but DON'T filter)
if len(cancels) > 0:
    cancels = cancels.join(
        orders_lookup,
        on="order_reference_number",
        how="left",
    )
    # NOT filtered

# Join E messages (but DON'T filter)
if len(executions) > 0:
    executions = executions.join(
        orders_lookup,
        on="order_reference_number",
        how="left",
    )
    # NOT filtered

# Join C messages (but DON'T filter)
if executions_c is not None and len(executions_c) > 0:
    executions_c = executions_c.join(
        orders_lookup,
        on="order_reference_number",
        how="left",
    )
    # NOT filtered

# U (Replace) messages: DON'T filter on join!
# Replace chains can be: A → U → U → U (order created by U, then replaced again)
# The original_order_reference_number may reference an order created by a previous U,
# not an A/F add. We look up the original order from the pool during reconstruction.
if replaces is not None and len(replaces) > 0:
    # Join to get original_side for orders that started with A/F
    replaces = replaces.join(
        orders_lookup.rename(
            {
                "order_reference_number": "original_order_reference_number",
                "price": "original_price",
                "buy_sell_indicator": "original_side",
                "original_shares": "replaced_shares",
            }
        ),
        on="original_order_reference_number",
        how="left",
    )
    # Note: Do NOT filter on original_price.is_not_null() - many U messages
    # reference orders created by other U messages, not A/F. The side is
    # looked up from the order pool during reconstruction.

print(f"\nSymbol: {SYMBOL}")
print(f"Trading Date: {TRADING_DATE}")
print("\nMessage counts:")
print(f"  Add orders (A+F): {len(add_orders):,}")
print(f"  Deletes (D): {len(deletes):,}")
print(f"  Cancels (X): {len(cancels):,}")
print(f"  Executions (E): {len(executions):,}")
print(f"  Replaces (U): {len(replaces) if replaces is not None else 0:,}")
print(f"  Trades (P): {len(trades):,}")

add_orders.head(5)

# A handful of far-from-market and market-peg orders sit at sentinel prices
# (near $0 or ~$200k), so the raw min-max spans an implausibly wide range. The
# 1st-99th percentile shows where displayed liquidity actually rests.
_p = add_orders["price"]
print(f"Price range (min-max):    ${_p.min():,.2f} - ${_p.max():,.2f}")
print(f"Central range (1st-99th): ${_p.quantile(0.01):,.2f} - ${_p.quantile(0.99):,.2f}")

# Use configured time window (defaults to full RTH: 09:30-16:00)
# Note: We load ALL messages from start of day up to end_time because
# delete/execute messages reference orders placed hours earlier in pre-market.
start_time = datetime.strptime(f"{TRADING_DATE} {START_TIME}", "%Y-%m-%d %H:%M:%S")
end_time = datetime.strptime(f"{TRADING_DATE} {END_TIME}", "%Y-%m-%d %H:%M:%S")

# CRITICAL: Include ALL messages from start of day up to end_time
# Delete/Execute messages at 9:31 reference Add orders from 4:00 AM pre-market.
# Without the full message history, we can't track remaining shares correctly.
#
# The order pool is built incrementally:
# - 4:00 AM: Add order 12345, 500 shares
# - 5:00 AM: Execute 200 shares (remaining = 300)
# - 9:31 AM: Delete order 12345 (remove 300 shares, not 500!)
#
# Snapshots are only generated after start_time.
add_all = add_orders.filter(pl.col("timestamp") <= end_time)
del_all = deletes.filter(pl.col("timestamp") <= end_time)
can_all = cancels.filter(pl.col("timestamp") <= end_time)
exec_all = executions.filter(pl.col("timestamp") <= end_time)
exec_c_all = (
    executions_c.filter(pl.col("timestamp") <= end_time) if executions_c is not None else None
)
rep_all = replaces.filter(pl.col("timestamp") <= end_time) if replaces is not None else None

print(f"Messages for LOB reconstruction (up to {end_time.time()}):")
print(f"  Add orders: {len(add_all):,}")
print(f"  Deletes: {len(del_all):,}")
print(f"  Cancels: {len(can_all):,}")
print(f"  Executions (E): {len(exec_all):,}")
print(f"  Executions (C): {len(exec_c_all) if exec_c_all is not None else 0:,}")
print(f"  Replaces: {len(rep_all) if rep_all is not None else 0:,}")

# Reconstruct LOB with OFI using Numba-accelerated function
# This computes Order Flow Imbalance during the reconstruction pass:
# OFI = (Bid Adds - Bid Removes) - (Ask Adds - Ask Removes)
lob = reconstruct_lob_with_ofi(
    add_all,
    del_all,
    can_all,
    exec_all,
    executions_c=exec_c_all,
    replaces=rep_all,
    n_levels=5,
    snapshot_freq="1s",
)

# Filter to RTH snapshots only (reconstruction processes all messages from start of day)
if len(lob) > 0:
    lob = lob.filter(pl.col("timestamp") >= start_time)

assert len(lob) > 0, (
    f"LOB reconstruction returned 0 snapshots for {SYMBOL} on {TRADING_DATE}. "
    "Check that the trading date has parsed ITCH messages on disk."
)

print(f"LOB snapshots: {len(lob):,}")

lob.head()

valid_count = (lob["spread"] > 0).sum()
crossed_count = (lob["spread"] < 0).sum()
print(f"Valid spreads (spread > 0): {valid_count:,} ({valid_count / len(lob) * 100:.1f}%)")
print(f"Crossed quotes (spread < 0): {crossed_count:,} ({crossed_count / len(lob) * 100:.1f}%)")

lob.select(
    pl.col("ofi").mean().alias("mean"),
    pl.col("ofi").std().alias("std"),
    pl.col("ofi").quantile(0.5).alias("median"),
    pl.col("ofi").min().alias("min"),
    pl.col("ofi").max().alias("max"),
)

# Create output directory for symbol
symbol_dir = OUTPUT_DIR / SYMBOL
symbol_dir.mkdir(parents=True, exist_ok=True)

lob_pd = lob.to_pandas().set_index("timestamp")
ofi_1m = (
    lob_pd[["ofi", "bid_size_0", "ask_size_0"]]
    .resample("1min")
    .agg({"ofi": "sum", "bid_size_0": "mean", "ask_size_0": "mean"})
)

fig, axes = plt.subplots(2, 1, figsize=(14, 8), sharex=True)

ax1 = axes[0]
# Per-minute OFI (signed). Positive values indicate net buying pressure
# within the minute, negative net selling. Per-minute is more readable
# than the cumulative trace because the opening cross dominates the
# cumulant's scale.
ofi_series = ofi_1m["ofi"].fillna(0)
ofi_series.plot(ax=ax1, color="#1E3A5F", linewidth=1.0, label="1-min OFI")
ax1.axhline(0, color="gray", linestyle="--", linewidth=0.5)
# Y-range from the 1st/99th percentiles to keep extreme minutes from
# squashing the rest of the session.
_vals = ofi_series.to_numpy()
if len(_vals):
    _lo, _hi = np.nanpercentile(_vals, [1, 99])
    _pad = max(abs(_lo), abs(_hi)) * 0.15
    ax1.set_ylim(_lo - _pad, _hi + _pad)
ax1.set_title(f"{SYMBOL} per-minute order-flow imbalance (signed)")
ax1.set_ylabel("OFI (shares per minute)")
ax1.legend(loc="upper right")

ax2 = axes[1]
ofi_1m["bid_size_0"].plot(ax=ax2, label="Bid Depth (L0)", color="green", alpha=0.7)
ofi_1m["ask_size_0"].plot(ax=ax2, label="Ask Depth (L0)", color="red", alpha=0.7)
ax2.set_title(f"{SYMBOL} top-of-book depth — bid and ask sizes evolve through the session")
ax2.set_ylabel("Shares")
ax2.legend()

plt.tight_layout()
plt.show()

output_file = symbol_dir / "lob_snapshots.parquet"
lob.write_parquet(output_file)
print(f"Saved LOB snapshots to: {output_file}")
print(f"Shape: {lob.shape}")