"""A teaching lab using LOBSTER's public AAPL sample.

Python 3.10+, numpy and pandas. This is a counterfactual quote experiment,
not an exchange replay: historical executions do not prove our order fills.
"""

from pathlib import Path
import argparse
import hashlib
import io
import json
import sqlite3
from urllib.request import urlopen
from zipfile import ZipFile

import numpy as np
import pandas as pd


SAMPLE_URL = (
    "https://php.lobsterdata.com/info/sample/"
    "LOBSTER_SampleFile_AAPL_2012-06-21_1.zip"
)
STEM = "AAPL_2012-06-21_34200000_57600000"
MESSAGE_COLUMNS = ["time", "event", "order_id", "size", "price", "direction"]
BOOK_COLUMNS = ["ask", "ask_size", "bid", "bid_size"]


def load_sample(data_dir: Path) -> pd.DataFrame:
    """Use two existing CSVs or download the provider's sample archive."""
    data_dir.mkdir(parents=True, exist_ok=True)
    message_path = data_dir / f"{STEM}_message_1.csv"
    book_path = data_dir / f"{STEM}_orderbook_1.csv"
    if message_path.exists() and book_path.exists():
        messages = pd.read_csv(message_path, header=None, names=MESSAGE_COLUMNS)
        book = pd.read_csv(book_path, header=None, names=BOOK_COLUMNS)
    else:
        archive_path = data_dir / "aapl-sample.zip"
        if not archive_path.exists():
            with urlopen(SAMPLE_URL, timeout=60) as response:
                archive_path.write_bytes(response.read())
        content = archive_path.read_bytes()
        (data_dir / "sample.sha256").write_text(hashlib.sha256(content).hexdigest() + "\n")
        with ZipFile(io.BytesIO(content)) as archive:
            messages = pd.read_csv(
                archive.open(f"{STEM}_message_1.csv"), header=None, names=MESSAGE_COLUMNS
            )
            book = pd.read_csv(
                archive.open(f"{STEM}_orderbook_1.csv"), header=None, names=BOOK_COLUMNS
            )
    if len(messages) != len(book):
        raise ValueError("Message and order-book row counts differ.")
    data = pd.concat([messages, book], axis=1)
    data.insert(0, "sequence", np.arange(len(data)))
    if not data["time"].is_monotonic_increasing:
        raise ValueError("Events are not in timestamp order.")
    # This lab requires an uninterrupted, valid, single-venue continuous session.
    if data["event"].eq(7).any():
        raise ValueError("Halt messages require explicit session segmentation.")
    valid = (
        (data["bid"] > 0) & (data["ask"] > data["bid"])
        & (data["ask"] < 9999999999)
        & (data["bid_size"] > 0) & (data["ask_size"] > 0)
    )
    if not valid.all():
        raise ValueError("Invalid or non-continuous book: segment before analysis.")
    # LOBSTER prices are dollars multiplied by 10,000, not cents.
    data[["price", "bid", "ask"]] = data[["price", "bid", "ask"]] / 10_000
    data["mid"] = (data["bid"] + data["ask"]) / 2
    data["spread"] = data["ask"] - data["bid"]
    data["imbalance"] = (
        (data["bid_size"] - data["ask_size"])
        / (data["bid_size"] + data["ask_size"])
    )
    return data


def passive_markouts(data: pd.DataFrame, horizon: float = 1.0) -> pd.DataFrame:
    """Value actual displayed executions from the passive participant's side."""
    times = data["time"].to_numpy()
    mids = data["mid"].to_numpy()
    targets = times + horizon
    # Last known midpoint at the target time; never use a later observation.
    future_index = np.searchsorted(times, targets, side="right") - 1
    eligible = data["event"].eq(4).to_numpy() & (targets <= times[-1])
    eligible[0] = False  # no pre-event snapshot for the first message
    row = np.flatnonzero(eligible)
    result = data.iloc[row][["sequence", "time", "size", "price", "direction"]].copy()
    # Direction describes the resting order: +1 is a passive buy, -1 a passive sell.
    side = result["direction"].to_numpy()
    pre_mid = mids[row - 1]
    future_mid = mids[future_index[row]]
    trade_price = result["price"].to_numpy()
    result["initial_edge"] = side * (pre_mid - trade_price)
    result["subsequent_move"] = side * (future_mid - pre_mid)
    result["markout"] = side * (future_mid - trade_price)
    if not np.allclose(result["markout"], result["initial_edge"] + result["subsequent_move"]):
        raise AssertionError("Markout accounting does not reconcile.")
    return result


def quote_experiment(
    data: pd.DataFrame,
    skew: float = 0.0,
    fill_probability: float = 0.5,
    fee: float = 0.003,
    limit: int = 100,
    seed: int = 42,
) -> dict:
    """One-share hypothetical quotes, instant updates and an assumed fill coin.

    Historical trade price reaching our quote creates an opportunity, not a
    proven fill. The arbitrary coin stands in for execution uncertainty.
    We ignore market impact, latency, hidden liquidity and queue reconstruction.
    """
    if not 0 <= fill_probability <= 1 or fee < 0 or limit < 1:
        raise ValueError("Invalid execution parameters.")
    values = data[["event", "price", "direction", "bid", "ask", "mid"]].to_numpy()
    uniforms = np.random.default_rng(seed).random(len(data))
    inventory, cash, fees, fills, opportunities, max_inventory = 0, 0.0, 0.0, 0, 0, 0
    half_width = 0.02  # dollars: nominal full spread is four cents
    for i in range(1, len(values)):
        event, trade_price, direction, _, _, _ = values[i]
        if event != 4:
            continue
        _, _, _, best_bid, best_ask, mid = values[i - 1]
        # Quote using the PREVIOUS row. Row i is the book AFTER event i.
        centre = mid - skew * inventory
        bid = np.floor((centre - half_width) * 100 + 1e-8) / 100
        ask = np.ceil((centre + half_width) * 100 - 1e-8) / 100
        # Clamp to passive prices, on the assumed $0.01 grid for this sample.
        bid = min(bid, best_ask - 0.01)
        ask = max(ask, best_bid + 0.01)
        buy = direction == 1 and trade_price <= bid + 1e-8 and inventory < limit
        sell = direction == -1 and trade_price >= ask - 1e-8 and inventory > -limit
        if not (buy or sell):
            continue
        opportunities += 1
        if uniforms[i] >= fill_probability:
            continue
        side, price = (1, bid) if buy else (-1, ask)
        # One share per selected execution message; no reuse of event volume.
        inventory += side
        cash -= side * price + fee
        fees += fee
        fills += 1
        max_inventory = max(max_inventory, abs(inventory))
    final_mid = float(data["mid"].iloc[-1])
    return {
        "skew_dollars_per_share": skew, "assumed_fill_probability": fill_probability,
        "opportunities": opportunities, "fills": fills,
        "ending_inventory": inventory, "max_abs_inventory": max_inventory,
        "fees": round(fees, 3), "gross_marked_pnl": round(cash + inventory * final_mid + fees, 3),
        "net_marked_pnl": round(cash + inventory * final_mid, 3),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data-dir", type=Path, default=Path("data/lobster"))
    parser.add_argument("--output-dir", type=Path, default=Path("market-making-output"))
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    data = load_sample(args.data_dir)
    marks = passive_markouts(data)
    # Keep event sequence: timestamps need not be unique.
    with sqlite3.connect(args.output_dir / "market_making.sqlite") as connection:
        data.to_sql("book_events", connection, index=False, if_exists="replace")
        marks.to_sql("passive_markouts", connection, index=False, if_exists="replace")
        connection.execute("CREATE UNIQUE INDEX event_sequence ON book_events(sequence)")
        connection.execute("CREATE INDEX event_time ON book_events(time)")
    marks.to_csv(args.output_dir / "passive_markouts.csv", index=False)
    results = pd.DataFrame([
        quote_experiment(data, skew=skew, fill_probability=p)
        for p in (0.25, 0.5, 1.0) for skew in (0.0, 0.001)
    ])
    results.to_csv(args.output_dir / "quote_experiments.csv", index=False)
    weights = marks["size"].to_numpy()
    summary = {
        "ticker": "AAPL", "date": "2012-06-21", "events": len(data),
        "first_bid": float(data["bid"].iloc[0]), "first_ask": float(data["ask"].iloc[0]),
        "median_event_spread": round(float(data["spread"].median()), 4),
        "displayed_executions_with_1s_horizon": len(marks),
        "volume_weighted_initial_edge": round(float(np.average(marks["initial_edge"], weights=weights)), 6),
        "volume_weighted_subsequent_move": round(float(np.average(marks["subsequent_move"], weights=weights)), 6),
        "volume_weighted_1s_markout": round(float(np.average(marks["markout"], weights=weights)), 6),
    }
    (args.output_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n")
    print(json.dumps(summary, indent=2))
    print(results.to_string(index=False))


if __name__ == "__main__":
    main()