#!/usr/bin/env python3
"""Fill-at-touch gap on an exchange's public data. Deterministic, single script. See fillgap-README.md.
A resting order fills under the proxy when volume at its price clears the queue ahead, or when a
trade prints through its price (the level was swept). The naive rule fills on any print at or through."""
import os, sys, glob, zipfile, argparse
import numpy as np, pandas as pd

HERE = os.path.dirname(os.path.abspath(__file__))
DATA = os.path.join(HERE, "data")
SIZES = {"BTCUSDT": 0.01, "ETHUSDT": 0.01, "SOLUSDT": 1.0}   # own order size, base-asset units (about $35 to $700)
TICKS = {"BTCUSDT": 0.1, "ETHUSDT": 0.01, "SOLUSDT": 0.001}  # USDT-M futures price tick, per symbol
SIZE, TICK = SIZES["BTCUSDT"], TICKS["BTCUSDT"]               # overwritten from --sym / --size / --tick in main()
TARGET_BPS, STOP_BPS, MAX_HOLD_MS = 2.0, 5.0, 60_000
FEE_BPS_CASES = [0.0, 1.0]
VOL_WIN_MS, BAR_MS = 300_000, 100          # 5-min realized range window; P&L path bar


def load_day(sym, day):
    """bookTicker + aggTrades for one UTC day. Returns (book df, trades df)."""
    bz = f"{DATA}/{sym}-bookTicker-{day}.zip"
    tz = f"{DATA}/{sym}-aggTrades-{day}.zip"
    with zipfile.ZipFile(bz) as z:
        book = pd.read_csv(z.open(z.namelist()[0]),
                           usecols=["best_bid_price", "best_bid_qty", "best_ask_price",
                                    "best_ask_qty", "transaction_time"],
                           dtype={"best_bid_price": "f8", "best_bid_qty": "f8",
                                  "best_ask_price": "f8", "best_ask_qty": "f8",
                                  "transaction_time": "i8"})
    with zipfile.ZipFile(tz) as z:
        tr = pd.read_csv(z.open(z.namelist()[0]),
                         usecols=["price", "quantity", "transact_time", "is_buyer_maker"])
    # A TICK coarser than the venue's merges neighbouring price levels, which silently
    # reclassifies sweeps as volume "at our price". Fail loudly instead of reporting it.
    for c in ("best_bid_price", "best_ask_price"):
        v = book[c].to_numpy()
        # Share, not any(): a wrong TICK puts most rows off the grid (0.1 on a 0.01-tick
        # symbol leaves 9 endings in 10 off it), while the archive itself carries the odd
        # malformed price — 2024-03-21 has exactly one, 65285.56 in 18.9M rows.
        off = (np.abs(v - np.rint(v / TICK) * TICK) > TICK * 1e-3).mean()
        if off > 0.01:
            sys.exit(f"{sym} {day}: {off:.1%} of {c} is not a multiple of TICK={TICK}. "
                     f"Pass --tick with the venue's price filter for {sym}.")
    book = book.sort_values("transaction_time", kind="stable").reset_index(drop=True)
    tr = tr.sort_values("transact_time", kind="stable").reset_index(drop=True)
    return book, tr


def first_through(tt_all, tp_all, t_join, t_end, Pk, sgn):
    """Per join: time of the first hit trade strictly through P (below a bid, above an ask) in
    (t_join, t_end], else -1. tt_all/tp_all are the hit trades in time order."""
    out = np.full(len(Pk), -1, dtype=np.int64)
    a = np.searchsorted(tt_all, t_join, "right")
    b = np.searchsorted(tt_all, t_end, "right")
    sp = sgn * tp_all                                              # "through" is always sp < sgn*P
    for j in np.flatnonzero(b > a):
        m = sp[a[j]:b[j]] < sgn * Pk[j]
        if m.any():
            out[j] = tt_all[a[j] + int(m.argmax())]
    return out


def joins_for_side(book, tr, side):
    """One row per touch-level run (= one simulated join). Vectorised fill test."""
    px = book[f"best_{side}_price"].to_numpy()
    qty = book[f"best_{side}_qty"].to_numpy()
    t = book["transaction_time"].to_numpy()
    chg = np.flatnonzero(np.r_[True, px[1:] != px[:-1]])          # run starts
    t_join, P, Q = t[chg], px[chg], qty[chg]
    t_end = np.r_[t[chg[1:]], t[-1]]                               # touch moves away
    # trades that hit this side: sells hit the bid (is_buyer_maker=True), buys lift the ask
    hit = tr["is_buyer_maker"].to_numpy() == (side == "bid")
    tp = np.rint(tr["price"].to_numpy()[hit] / TICK).astype(np.int64)
    tt = tr["transact_time"].to_numpy()[hit]
    tq = tr["quantity"].to_numpy()[hit]
    Pk = np.rint(P / TICK).astype(np.int64)
    # prints THROUGH our price: a sell below a bid at P (or a buy above an ask at P) can only execute
    # once every order resting at P, ours included, has filled. Fills both rules from the join.
    t_thr = first_through(tt, tp, t_join, t_end, Pk, 1 if side == "bid" else -1)
    t0 = min(t[0], tt[0]) if len(tt) else t[0]
    key = tp * (1 << 32) + (tt - t0)                               # sort by (price, time)
    o = np.argsort(key, kind="stable"); key, tt, tq = key[o], tt[o], tq[o]
    cs = np.r_[0.0, np.cumsum(tq)]
    lo = np.searchsorted(key, Pk * (1 << 32) + (t_join - t0), "right")   # window (t_join, t_end]
    hi = np.searchsorted(key, Pk * (1 << 32) + (t_end - t0), "right")
    at = hi > lo                                                   # some print at exactly P
    k = np.searchsorted(cs, cs[lo] + Q + SIZE - 1e-12, "left")        # first trade index+1 reaching queue
    filled_at = at & (k <= hi) & (k > lo)
    tf_at = np.where(filled_at, tt[np.minimum(k, len(tt)) - 1], -1)
    tn_at = np.where(at, tt[np.minimum(lo, len(tt) - 1)], -1)          # first print at P
    first = lambda a, b: np.where((a >= 0) & (b >= 0), np.minimum(a, b), np.maximum(a, b))
    tn, tf = first(tn_at, t_thr), first(tf_at, t_thr)              # earliest of at-P and through
    swept = (t_thr >= 0) & ((tf_at < 0) | (t_thr <= tf_at))          # the through print is what filled us
    return pd.DataFrame({"side": side, "t_join": t_join, "t_end": t_end, "price": P, "queue": Q,
                         "touched": tn >= 0, "naive": tn >= 0, "proxy": tf >= 0,
                         "through": swept, "at_only": filled_at & ~swept,
                         "t_fill_proxy": tf, "t_fill_naive": tn})


def vol_buckets(book, joins):
    """Trailing 5-min realized range of mid (bps) at each join time."""
    mid = (book["best_bid_price"].to_numpy() + book["best_ask_price"].to_numpy()) / 2
    t = book["transaction_time"].to_numpy()
    b = (t - t[0]) // 60_000                                            # 1-min bins
    df = pd.DataFrame({"b": b, "m": mid}).groupby("b")["m"].agg(["max", "min"])
    df = df.reindex(range(int(b.max()) + 1)).ffill()
    r = (df["max"].rolling(5, min_periods=1).max() - df["min"].rolling(5, min_periods=1).min()) \
        / df["min"] * 1e4
    jb = np.clip((joins["t_join"].to_numpy() - t[0]) // 60_000 - 1, 0, len(r) - 1)   # prior window
    return r.to_numpy()[jb]


def toy_pnl(book, joins, rule):
    """Toy MR maker: after fill at P, exit at +TARGET / -STOP / 60s at mid. bps per trade."""
    mid = (book["best_bid_price"].to_numpy() + book["best_ask_price"].to_numpy()) / 2
    t = book["transaction_time"].to_numpy()
    b = (t - t[0]) // BAR_MS
    bars = pd.DataFrame({"b": b, "m": mid}).groupby("b")["m"].agg(["max", "min", "last"])
    bars = bars.reindex(range(int(b.max()) + 1)).ffill()
    hi, lo, last = bars["max"].to_numpy(), bars["min"].to_numpy(), bars["last"].to_numpy()
    f = joins[joins[rule]]
    tf, P = f[f"t_fill_{rule}"].to_numpy(), f["price"].to_numpy()
    sgn = np.where(f["side"].to_numpy() == "bid", 1.0, -1.0)          # bid fill = long
    b0 = (tf - t[0]) // BAR_MS + 1
    n = MAX_HOLD_MS // BAR_MS
    out, oc = np.empty(len(f)), np.empty(len(f), dtype=object)
    for i in range(len(f)):
        s, e = int(b0[i]), min(int(b0[i]) + n, len(hi))
        if s >= len(hi):
            out[i], oc[i] = 0.0, "timeout"; continue
        tgt, stp = P[i] * (1 + sgn[i] * TARGET_BPS / 1e4), P[i] * (1 - sgn[i] * STOP_BPS / 1e4)
        h, l = hi[s:e], lo[s:e]
        st = np.flatnonzero((l <= stp) if sgn[i] > 0 else (h >= stp))
        tg = np.flatnonzero((h >= tgt) if sgn[i] > 0 else (l <= tgt))
        js, jt = (st[0] if len(st) else n), (tg[0] if len(tg) else n)
        if js <= jt and js < n: out[i], oc[i] = -STOP_BPS, "stop"      # stop wins ties (conservative)
        elif jt < n: out[i], oc[i] = TARGET_BPS, "target"
        else: out[i], oc[i] = sgn[i] * (last[e - 1] - P[i]) / P[i] * 1e4, "timeout"
    return pd.DataFrame({"day": f["t_join"].to_numpy() // 86_400_000, "bps": out, "price": P, "outcome": oc,
                         "kind": np.where(f["through"].to_numpy(), "through", "at_p"),
                         "proxy_filled": f["proxy"].to_numpy()})      # for the naive-only row


def run_day(sym, day):
    book, tr = load_day(sym, day)
    j = pd.concat([joins_for_side(book, tr, "bid"), joins_for_side(book, tr, "ask")], ignore_index=True)
    j["day"] = day
    j["hour"] = (j["t_join"] % 86_400_000) // 3_600_000
    j["vol_bps"] = vol_buckets(book, j)
    pnl = {r: toy_pnl(book, j, r) for r in ("naive", "proxy")}
    return j, pnl


def rate_table(j, by):
    g = j.groupby(by)
    df = pd.DataFrame({"joins": g.size(), "naive_fill": g["naive"].mean(), "proxy_fill": g["proxy"].mean(),
                       "proxy_via_through": g["through"].mean(), "proxy_at_p_only": g["at_only"].mean()})
    df["gap_pp"] = (df["naive_fill"] - df["proxy_fill"]) * 100
    return df


def plot_charts(by_hour, pnl_day, sym, sfx):
    """The two charts, drawn dark so they sit on the page. Inputs are the by-hour and per-day P&L tables."""
    import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
    BG, FG, GRID = "#0E0805", "#D4D4D4", "#3A2418"
    plt.rcParams.update({"figure.facecolor": BG, "axes.facecolor": BG, "savefig.facecolor": BG, "axes.edgecolor": GRID,
                         "axes.labelcolor": FG, "text.color": FG, "xtick.color": FG, "ytick.color": FG, "grid.color": GRID,
                         "legend.frameon": False, "legend.labelcolor": FG, "font.size": 11})
    fig, ax = plt.subplots(figsize=(9, 4))
    ax.plot(by_hour.index, by_hour["naive_fill"] * 100, marker="o", color="#E66B2C", label="naive (touch = fill)")
    ax.plot(by_hour.index, by_hour["proxy_fill"] * 100, marker="o", color="#4A8FBF", label="queue proxy")
    ax.set_xlabel("hour (UTC)"); ax.set_ylabel("fill rate %"); ax.set_title(f"{sym}: fill rate at touch by hour", color="#fff")
    ax.legend(); ax.grid(alpha=.5)
    fig.tight_layout(); fig.savefig(f"{HERE}/fill_by_hour{sfx}.png", dpi=130); plt.close(fig)
    fig, ax = plt.subplots(figsize=(9, 4))
    style = {"naive_fee0bps": ("naive, no fee", "#FF8B4D", "--"), "naive_fee1bps": ("naive, 1 bp fee", "#E66B2C", "-"),
             "proxy_fee0bps": ("queue proxy, no fee", "#7FB2DD", "--"), "proxy_fee1bps": ("queue proxy, 1 bp fee", "#4A8FBF", "-")}
    for c in pnl_day.columns:
        lab, col, ls = style.get(c, (c, FG, "-"))
        ax.plot(pnl_day.index[:-1], pnl_day[c].iloc[:-1].cumsum(), marker="o", color=col, ls=ls, label=lab)
    ax.set_ylabel("cumulative P&L (USD)"); ax.set_title("Toy maker: naive vs proxy fill rule", color="#fff"); ax.legend(); ax.grid(alpha=.5)
    ax.tick_params(axis="x", rotation=45); fig.tight_layout(); fig.savefig(f"{HERE}/cum_pnl{sfx}.png", dpi=130); plt.close(fig)


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--sym", default="BTCUSDT"); ap.add_argument("--days", default="", help="comma list, default all")
    ap.add_argument("--tick", type=float, default=None, help="price tick; default TICKS[sym]")
    ap.add_argument("--size", type=float, default=None, help="own order size in base units; default SIZES[sym], else 0.01")
    ap.add_argument("--export-joins", action="store_true", help="also write joins{-sym}.csv, the join log for `bt2live fill-gap --latency-ms N`")
    a = ap.parse_args(); sfx = "" if a.sym == "BTCUSDT" else f"-{a.sym.lower()}"   # ETH etc. get suffixed outputs
    base = a.sym.removesuffix("USDT")          # SIZE is in the base asset, not always BTC
    global TICK, SIZE
    SIZE = a.size if a.size is not None else SIZES.get(a.sym, SIZE)
    if a.tick is not None: TICK = a.tick
    elif a.sym in TICKS: TICK = TICKS[a.sym]
    else: sys.exit(f"no tick known for {a.sym}; pass --tick with the venue's price filter")
    days = sorted(f.split("-bookTicker-")[1][:-4] for f in glob.glob(f"{DATA}/{a.sym}-bookTicker-*.zip"))
    if a.days: days = [d for d in days if d in a.days.split(",")]
    J, PN = [], {"naive": [], "proxy": []}
    for d in days:
        j, pnl = run_day(a.sym, d); J.append(j)
        for r in PN: PN[r].append(pnl[r])
        print(d, len(j), "joins; proxy fill", round(j["proxy"].mean(), 4), "naive", round(j["naive"].mean(), 4), file=sys.stderr)
    j = pd.concat(J, ignore_index=True)
    if a.export_joins:    # one row per simulated join. No "filled" column on purpose: with one, bt2live measures
        # against that 0 ms flag at every --latency-ms; without it, against its own touch rule at the delayed arrival
        pd.DataFrame({"ts": j["t_join"], "side": np.where(j["side"] == "bid", "buy", "sell"),
                      "price": (np.rint(j["price"] / TICK) * TICK).round(8), "size": SIZE,
                      "end_ts": j["t_end"]}).to_csv(f"{HERE}/joins{sfx}.csv", index=False)
    j["vol_tercile"] = pd.qcut(j["vol_bps"], 3, labels=["low", "mid", "high"])
    overall = rate_table(j, "side"); overall.loc["all"] = rate_table(j.assign(k="all"), "k").loc["all"]
    ttf = (j.loc[j["proxy"], "t_fill_proxy"] - j.loc[j["proxy"], "t_join"])
    ttf_n = (j.loc[j["naive"], "t_fill_naive"] - j.loc[j["naive"], "t_join"])
    by_hour, by_vol, by_day = rate_table(j, "hour"), rate_table(j, "vol_tercile"), rate_table(j, "day")
    # toy P&L per day, USD for SIZE of the base asset, both fill rules x fee cases
    rows = []
    for r in PN:
        p = pd.concat(PN[r]);
        for fee in FEE_BPS_CASES:
            usd = (p["bps"] - fee) / 1e4 * p["price"] * SIZE
            rows.append(usd.groupby(p["day"]).sum().rename(f"{r}_fee{fee:g}bps"))
    pnl_day = pd.concat(rows, axis=1); pnl_day.index = pd.to_datetime(pnl_day.index, unit="D").strftime("%Y-%m-%d")
    pnl_day.loc["total"] = pnl_day.sum()
    def diag(p):
        return {"fills": len(p), "mean_bps_per_fill": p["bps"].mean(),
                **p["outcome"].value_counts(normalize=True).rename(lambda k: f"share_{k}").to_dict()}
    pn, pp = pd.concat(PN["naive"]), pd.concat(PN["proxy"])
    per_fill = pd.DataFrame({"naive": diag(pn), "proxy": diag(pp),
                             "proxy: level swept (through)": diag(pp[pp["kind"] == "through"]),
                             "proxy: queue cleared at P": diag(pp[pp["kind"] == "at_p"]),
                             "naive only (no proxy fill)": diag(pn[~pn["proxy_filled"]])}).T
    fmt = lambda df: df.to_markdown(floatfmt=".4f")
    with open(f"{HERE}/results{sfx}.md", "w") as f:
        f.write(f"# Results: {a.sym} (USDT-M perpetual futures, public exchange archive), {days[0]} .. {days[-1]}\n\n")
        f.write(f"Joins simulated: {len(j):,} ({len(days)} days, both sides). Own size {SIZE} {base}.\n\n")
        f.write("## Fill rate overall\n\n" + fmt(overall) + "\n\n")
        f.write(f"Median time-to-fill (proxy): {ttf.median():.0f} ms; (naive, first touch): {ttf_n.median():.0f} ms. "
                f"Proxy p25/p75: {ttf.quantile(.25):.0f}/{ttf.quantile(.75):.0f} ms.\n\n")
        f.write("## By hour of day (UTC)\n\n" + fmt(by_hour) + "\n\n")
        f.write("## By volatility tercile (trailing 5-min realized range of mid, bps)\n\n")
        f.write(j.groupby("vol_tercile", observed=True)["vol_bps"].agg(["min", "max"]).to_markdown(floatfmt=".1f") + "\n\n" + fmt(by_vol) + "\n\n")
        f.write("## By day\n\n" + fmt(by_day) + "\n\n")
        f.write(f"## Toy mean-reversion maker P&L per day (USD, {SIZE} {base} per fill)\n\n" + pnl_day.to_markdown(floatfmt=".2f") + "\n\n")
        f.write("Per-fill diagnostics (before fees):\n\n" + per_fill.to_markdown(floatfmt=".3f") + "\n\n")
        f.write(f"![fill by hour](fill_by_hour{sfx}.png)\n\n![cum pnl](cum_pnl{sfx}.png)\n")
    plot_charts(by_hour, pnl_day, a.sym, sfx)
    o = overall.loc["all"]
    with open(f"{HERE}/summary{sfx}.txt", "w") as f:
        f.write(f"1. Naive fill rate at touch (any print at or through our price): {o['naive_fill']*100:.1f}%\n")
        f.write(f"2. Queue-proxy fill rate at touch: {o['proxy_fill']*100:.1f}%  (gap {o['gap_pp']:.1f} pp; proxy/naive = {o['proxy_fill']/o['naive_fill']:.2f}; "
                f"of joins, {o['proxy_via_through']*100:.1f}% fill via a print through the level, {o['proxy_at_p_only']*100:.1f}% via volume at the level only)\n")
        f.write(f"3. Median time-to-fill: proxy {ttf.median()/1000:.2f}s vs naive {ttf_n.median()/1000:.2f}s\n")
        f.write(f"4. Proxy fill rate low-vol vs high-vol tercile: {by_vol.loc['low','proxy_fill']*100:.1f}% vs {by_vol.loc['high','proxy_fill']*100:.1f}%\n")
        f.write(f"5. Toy MR maker {len(days)}-day P&L, 1 bps fee: naive {pnl_day.loc['total','naive_fee1bps']:+.0f} USD vs proxy {pnl_day.loc['total','proxy_fee1bps']:+.0f} USD (0 fee: {pnl_day.loc['total','naive_fee0bps']:+.0f} vs {pnl_day.loc['total','proxy_fee0bps']:+.0f})\n")
    print(open(f"{HERE}/summary{sfx}.txt").read())


if __name__ == "__main__":
    main()