"""Where and when gold moves violently — burst map of the raw bars.

Usage:
    python backtest/research/burst_map.py --parquet PATH --bar-minutes N \
        [--k 3,4,5] [--book BOOK.csv] [--within-bars 12] \
        [--start YYYY-MM-DD --end YYYY-MM-DD] [--out DIR]

A "burst" is one bar whose range is >= K x the PRIOR bar's Wilder ATR(14) —
the same definition burst_profile.py uses (that module computes it inline
inside burst_table() and exposes no reusable helper, so the formula is
replicated here bar-for-bar: TR = max(H-L, |H-Cprev|, |L-Cprev|), ATR =
TR.ewm(alpha=1/14, min_periods=14, adjust=False).mean(), ratio =
(H-L) / ATR.shift(1)).

Reports:
  * burst frequency by (weekday, UTC bar-CLOSE hour), plus mean |1 h| and
    |4 h| absolute moves in $ and in ATR units on the same grid;
  * post-burst drift in the burst's own direction over 6/12/24/48 bars,
    overall and by session bucket — does a violent bar continue or snap back;
  * the largest trailing-60-minute close-to-close moves, de-overlapped;
  * if a trade book is given: how trades that were decided just after a burst,
    or that were OPEN when one printed, performed against the rest.

Every ATR-unit number here is normalised by the PRIOR bar's ATR — the same
denominator as the burst test, so a "4x burst that ran another 4 ATR" reads
directly. All hours are UTC and are the bar's CLOSE hour (label + one bar),
matching the clock the live gates evaluate.

Analysis tooling only — never imported by the live code.
"""
from __future__ import annotations

import argparse
import math
import os

import numpy as np
import pandas as pd

try:  # package import (backtest.research.burst_map)
    from backtest.research.book_report import base_metrics, load_book, session_of
except ImportError:  # run directly as a script from anywhere
    from book_report import base_metrics, load_book, session_of

SLOT_UNKNOWN = ""


# ─────────────────────────── bar primitives ─────────────────────────

def load_bars(parquet: str, start: str | None = None,
              end: str | None = None) -> pd.DataFrame:
    """Read a cached OHLCV parquet, de-duplicated and oldest-first."""
    df = pd.read_parquet(parquet).sort_index()
    df = df[~df.index.duplicated(keep="last")]
    if start:
        df = df[df.index >= pd.Timestamp(start, tz="UTC")]
    if end:
        df = df[df.index <= pd.Timestamp(end, tz="UTC")]
    return df


def wilder_atr(df: pd.DataFrame, period: int = 14) -> pd.Series:
    """Wilder ATR — identical to the inline formula in burst_profile.py."""
    h, l, c = df["High"], df["Low"], df["Close"]
    tr = pd.concat([h - l, (h - c.shift()).abs(), (l - c.shift()).abs()],
                   axis=1).max(axis=1)
    return tr.ewm(alpha=1 / period, min_periods=period, adjust=False).mean()


def burst_ratio(df: pd.DataFrame, period: int = 14) -> pd.Series:
    """Bar range expressed in PRIOR-bar ATRs."""
    return (df["High"] - df["Low"]) / wilder_atr(df, period).shift(1)


def burst_flags(df: pd.DataFrame, k: float, period: int = 14) -> pd.Series:
    """True where the bar's range is >= k x the prior bar's Wilder ATR(14)."""
    return (burst_ratio(df, period) >= k).fillna(False)


def burst_dir(df: pd.DataFrame) -> pd.Series:
    """+1 up bar, −1 down bar, 0 doji (sign of Close − Open)."""
    return np.sign(df["Close"] - df["Open"]).astype(float)


def close_times(df: pd.DataFrame, bar_minutes: int) -> pd.DatetimeIndex:
    """Bar CLOSE timestamps (the parquet index labels bar OPEN time)."""
    return df.index + pd.Timedelta(minutes=bar_minutes)


# ───────────────────── when does gold move? tables ──────────────────

def hour_weekday_table(df: pd.DataFrame, k: float, bar_minutes: int = 5,
                       ks: tuple[float, ...] | None = None) -> dict:
    """Burst frequency and absolute-move size on a (weekday, close-hour) grid.

    Returns a dict of DataFrames (several tables are wanted, one frame each):
      burst_pct  : weekday rows x close-hour cols, % of bars that are >= k x
                   bursts (one frame; per-k frames live in `burst_pct_by_k`)
      absmove    : (weekday, close-hour) rows; bars, mean |1 h move| and
                   |4 h move| in $ and in prior-bar ATR units
      by_hour    : the same stats collapsed to close-hour
      by_weekday : the same stats collapsed to weekday
    """
    ks = tuple(ks) if ks else (k,)
    ct = close_times(df, bar_minutes)
    hour = pd.Series(ct.hour, index=df.index)
    wd = pd.Series(ct.weekday, index=df.index)
    atr_prev = wilder_atr(df).shift(1)

    per_hour = max(1, int(round(60 / bar_minutes)))
    c = df["Close"]
    mv1 = (c - c.shift(per_hour)).abs()
    mv4 = (c - c.shift(4 * per_hour)).abs()
    stats = pd.DataFrame({
        "weekday": wd, "hour": hour,
        "abs1h_usd": mv1, "abs4h_usd": mv4,
        "abs1h_atr": mv1 / atr_prev, "abs4h_atr": mv4 / atr_prev,
    })
    for kk in ks:
        stats[f"burst_{kk:g}x"] = burst_flags(df, kk).astype(float) * 100.0

    agg = {"abs1h_usd": "mean", "abs4h_usd": "mean",
           "abs1h_atr": "mean", "abs4h_atr": "mean"}
    agg.update({f"burst_{kk:g}x": "mean" for kk in ks})

    grid = stats.groupby(["weekday", "hour"]).agg(agg)
    grid.insert(0, "bars", stats.groupby(["weekday", "hour"]).size())
    by_hour = stats.groupby("hour").agg(agg)
    by_hour.insert(0, "bars", stats.groupby("hour").size())
    by_wd = stats.groupby("weekday").agg(agg)
    by_wd.insert(0, "bars", stats.groupby("weekday").size())

    pct_by_k = {kk: grid[f"burst_{kk:g}x"].unstack("hour") for kk in ks}
    return {
        "burst_pct": pct_by_k[k if k in pct_by_k else ks[0]].round(3),
        "burst_pct_by_k": {kk: v.round(3) for kk, v in pct_by_k.items()},
        "absmove": grid.round(3),
        "by_hour": by_hour.round(3),
        "by_weekday": by_wd.round(3),
    }


def post_burst_drift(df: pd.DataFrame, k: float,
                     horizons: tuple[int, ...] = (6, 12, 24, 48),
                     bar_minutes: int = 5) -> pd.DataFrame:
    """Signed drift IN the burst's direction over the next h bars, in ATRs.

    A positive mean means violent bars continue; negative means they snap
    back. Split by the session bucket of the burst bar's close hour.
    """
    flags = burst_flags(df, k)
    d = burst_dir(df)
    atr_prev = wilder_atr(df).shift(1)
    ct = close_times(df, bar_minutes)
    bucket = pd.Series([session_of(h) for h in ct.hour], index=df.index)
    c = df["Close"]

    rows = []
    for h in horizons:
        fwd = (c.shift(-h) - c) * d / atr_prev
        sub = fwd[flags].dropna()
        buckets = bucket[flags].reindex(sub.index)
        rows.append({"horizon_bars": h, "bucket": "ALL", "n": int(len(sub)),
                     "mean_atr": round(float(sub.mean()), 4) if len(sub) else float("nan"),
                     "median_atr": round(float(sub.median()), 4) if len(sub) else float("nan"),
                     "pct_positive": round(float((sub > 0).mean() * 100), 2) if len(sub) else float("nan")})
        for b, s in sub.groupby(buckets):
            rows.append({"horizon_bars": h, "bucket": b, "n": int(len(s)),
                         "mean_atr": round(float(s.mean()), 4),
                         "median_atr": round(float(s.median()), 4),
                         "pct_positive": round(float((s > 0).mean() * 100), 2)})
    return pd.DataFrame(rows)


def slot_label(weekday: int, hour: int, day: int) -> str:
    """Time-slot label for a big move.

    NOTE: there is no economic-calendar file in this repo, so these are
    TIME SLOTS, not confirmed events — "payroll window" means "the clock
    slot US payrolls normally print in", nothing more.
    """
    if weekday == 4 and day <= 7 and hour in (12, 13):
        return "first-Friday 12:30 slot (payroll-report window)"
    if weekday == 2 and hour in (18, 19):
        return "Wednesday 18:00 slot (central-bank window)"
    if hour in (12, 13):
        return "12:30 slot (US data window)"
    if hour == 14:
        return "14:00 slot"
    if hour in (21, 22, 23):
        return "settlement/after-hours"
    if hour in (0, 1, 2):
        return "Asia open"
    return SLOT_UNKNOWN


def largest_moves(df: pd.DataFrame, n: int = 40, window_minutes: int = 60,
                  bar_minutes: int = 5) -> pd.DataFrame:
    """The n biggest trailing-`window_minutes` close-to-close moves.

    Overlapping windows are de-duplicated greedily: the largest move wins and
    anything whose close time is within +/- `window_minutes` of an already
    kept move is dropped, so one violent hour contributes one row.
    """
    bars = max(1, int(round(window_minutes / bar_minutes)))
    c = df["Close"]
    move = c - c.shift(bars)
    atr_prev = wilder_atr(df).shift(1)
    ct = close_times(df, bar_minutes)
    tbl = pd.DataFrame({
        "close_time": ct, "move_usd": move.to_numpy(),
        "move_atr": (move / atr_prev).to_numpy(),
    }).dropna()
    tbl["abs_usd"] = tbl["move_usd"].abs()
    tbl = tbl.sort_values("abs_usd", ascending=False)

    keep, taken = [], []
    gap = pd.Timedelta(minutes=window_minutes)
    for _, row in tbl.iterrows():
        t = row["close_time"]
        if any(abs(t - u) <= gap for u in taken):
            continue
        taken.append(t)
        keep.append(row)
        if len(keep) >= n:
            break
    out = pd.DataFrame(keep).reset_index(drop=True)
    if out.empty:
        return out
    out["weekday"] = out["close_time"].dt.weekday
    out["hour"] = out["close_time"].dt.hour
    out["direction"] = np.where(out["move_usd"] > 0, "UP", "DOWN")
    out["slot"] = [slot_label(w, h, t.day) for w, h, t
                   in zip(out["weekday"], out["hour"], out["close_time"])]
    out["move_usd"] = out["move_usd"].round(2)
    out["move_atr"] = out["move_atr"].round(2)
    return out[["close_time", "weekday", "hour", "direction", "move_usd",
                "move_atr", "slot"]]


# ───────────────────────── book <-> burst join ──────────────────────

def _stats(r: pd.Series) -> dict:
    m = base_metrics(r)
    return {"n": m["n"], "total_r": round(m["total_r"], 2),
            "mean_r": round(m["mean_r"], 3), "wr": round(m["wr"], 2),
            "pf": round(m["pf"], 2) if math.isfinite(m["pf"]) else float("inf")}


def join_book_to_bursts(book_df: pd.DataFrame, bars: pd.DataFrame, k: float,
                        within_bars: int = 12,
                        bar_minutes: int = 5) -> dict:
    """Split a trade book by its relationship to burst bars.

    (a) `after_burst`  : decision_time falls within `within_bars` bars after
        the most recent burst's close (bar distance is counted POSITIONALLY on
        the bar index, so weekend gaps cannot inflate it).
    (b) `burst_during` : at least one burst closed inside [entry, exit).
    (c) `burst_agrees` / `burst_opposes`: of (b), whether the FIRST burst that
        printed during the trade pushed with or against the trade direction.
    """
    ct = close_times(bars, bar_minutes)
    ct_arr = ct.to_numpy()
    flags = burst_flags(bars, k).to_numpy()
    bdir = burst_dir(bars).to_numpy()
    burst_pos = np.flatnonzero(flags)

    # Last burst position at or before each bar (−1 where none yet).
    last_burst = np.where(flags, np.arange(len(flags)), -1)
    last_burst = np.maximum.accumulate(last_burst)

    dec = book_df["decision_time"].astype("datetime64[ns, UTC]").to_numpy()
    ent = book_df["entry_time"].astype("datetime64[ns, UTC]").to_numpy()
    exi = book_df["exit_time"].astype("datetime64[ns, UTC]").to_numpy()

    # Position of the last bar that had CLOSED by decision_time.
    idx = np.searchsorted(ct_arr, dec, side="right") - 1
    valid = idx >= 0
    lb = np.where(valid, last_burst[np.clip(idx, 0, len(flags) - 1)], -1)
    bars_since = np.where((lb >= 0) & valid, idx - lb, np.iinfo(np.int32).max)
    after = pd.Series((bars_since >= 0) & (bars_since <= within_bars),
                      index=book_df.index)

    # Bursts printing inside the open trade.
    bct = ct_arr[burst_pos] if len(burst_pos) else np.array([], dtype=ct_arr.dtype)
    lo = np.searchsorted(bct, ent, side="left")
    hi = np.searchsorted(bct, exi, side="left")
    during = pd.Series(hi > lo, index=book_df.index)
    trade_sign = np.where(
        book_df["direction"].astype(str).str.upper().str.startswith("B"),
        1.0, -1.0)
    has = hi > lo
    first_dir = np.zeros(len(book_df), dtype=float)
    if len(burst_pos) and has.any():
        # hi > lo guarantees lo[has] indexes a real burst.
        first_dir[has] = bdir[burst_pos[lo[has]]]
    agrees = pd.Series(has & (first_dir * trade_sign > 0),
                       index=book_df.index)
    opposes = pd.Series(has & (first_dir * trade_sign < 0),
                        index=book_df.index)

    r = book_df["r_multiple"]
    out = {
        "k": k, "within_bars": within_bars,
        "bursts_in_window": int(len(burst_pos)),
        "after_burst": _stats(r[after]),
        "not_after_burst": _stats(r[~after]),
        "burst_during": _stats(r[during]),
        "no_burst_during": _stats(r[~during]),
        "burst_agrees": _stats(r[agrees]),
        "burst_opposes": _stats(r[opposes]),
        "median_bars_since_burst": float(np.median(
            bars_since[bars_since < np.iinfo(np.int32).max])) if (
            bars_since < np.iinfo(np.int32).max).any() else float("nan"),
    }
    return out


# ──────────────────────────── reporting ─────────────────────────────

def _fmt_stats(name: str, s: dict) -> str:
    pf = "inf" if math.isinf(s["pf"]) else f"{s['pf']:.2f}"
    return (f"    {name:<18} n={s['n']:<5} totalR={s['total_r']:>8.2f}  "
            f"meanR={s['mean_r']:>7.3f}  WR={s['wr']:>6.2f}%  PF={pf}")


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--parquet", required=True)
    ap.add_argument("--bar-minutes", type=int, required=True)
    ap.add_argument("--k", default="3,4,5",
                    help="burst thresholds in prior-bar ATRs (default 3,4,5)")
    ap.add_argument("--book", default=None)
    ap.add_argument("--book-bar-minutes", type=int, default=None,
                    help="bar minutes of the BOOK (default: --bar-minutes)")
    ap.add_argument("--within-bars", type=int, default=12)
    ap.add_argument("--start", default=None)
    ap.add_argument("--end", default=None)
    ap.add_argument("--n-moves", type=int, default=40)
    ap.add_argument("--out", default="backtest/research/burst_map_out")
    args = ap.parse_args()

    ks = [float(x) for x in args.k.split(",") if x.strip()]
    os.makedirs(args.out, exist_ok=True)
    bars = load_bars(args.parquet, args.start, args.end)
    bm = args.bar_minutes
    primary = ks[0]

    lines: list[str] = []
    add = lines.append
    add("=" * 78)
    add(f"BURST MAP  {os.path.basename(args.parquet)}")
    add("=" * 78)
    add(f"bars             : {len(bars)}  ({bars.index[0]} -> {bars.index[-1]}"
        f", labels = bar OPEN time)")
    add(f"bar minutes      : {bm}   hours below are bar CLOSE hours (UTC)")
    add(f"burst definition : range >= K x PRIOR bar Wilder ATR(14); "
        f"K = {', '.join(f'{k:g}' for k in ks)}")
    add("ATR units        : always normalised by the PRIOR bar's ATR")
    for k in ks:
        f = burst_flags(bars, k)
        add(f"  K={k:g}: {int(f.sum())} burst bars "
            f"({float(f.mean()) * 100:.3f}% of bars)")

    tabs = hour_weekday_table(bars, primary, bm, ks=tuple(ks))
    add("")
    add(f"── BURST % BY (weekday, close-hour), K={primary:g} ──")
    add(tabs["burst_pct"].fillna(0).round(2).to_string())
    add("")
    add("── MEAN ABSOLUTE MOVE BY CLOSE-HOUR (all weekdays) ──")
    add(tabs["by_hour"].to_string())
    add("")
    add("── MEAN ABSOLUTE MOVE BY WEEKDAY ──")
    add(tabs["by_weekday"].to_string())

    add("")
    add(f"── POST-BURST DRIFT (in the burst's own direction), K={primary:g} ──")
    drift = post_burst_drift(bars, primary, bar_minutes=bm)
    add(drift.to_string(index=False))

    add("")
    add(f"── LARGEST {args.n_moves} TRAILING-60-MIN MOVES (de-overlapped) ──")
    add("   slot labels are CLOCK SLOTS — no economic calendar exists in "
        "this repo, so no event is confirmed.")
    moves = largest_moves(bars, n=args.n_moves, bar_minutes=bm)
    add(moves.to_string(index=False))

    join_rows = []
    if args.book:
        bbm = args.book_bar_minutes or bm
        book = load_book(args.book, bbm)
        add("")
        add(f"── BOOK JOIN  {os.path.basename(args.book)} "
            f"({len(book)} trades, book bar = {bbm} min) ──")
        first, last = close_times(bars, bm)[0], close_times(bars, bm)[-1]
        outside = int(((book["decision_time"] < first)
                       | (book["exit_time"] > last)).sum())
        if outside:
            add(f"  !! {outside} trade(s) start before / end after the bar "
                f"window ({first} .. {last}) — their burst context is "
                f"partial. Widen --start/--end or use a longer parquet.")
        add("  CAUTION: group (b) is survivorship-shaped — a trade only sees "
            "a burst if it survived long enough to still be open, so its high "
            "PF is mostly 'winners last longer', not a burst edge.")
        for k in ks:
            j = join_book_to_bursts(book, bars, k, args.within_bars, bm)
            add(f"  K={k:g}  ({j['bursts_in_window']} bursts in window; "
                f"median bars since last burst = "
                f"{j['median_bars_since_burst']:.0f})")
            add(f"    (a) decided within {args.within_bars} bars of a burst")
            add(_fmt_stats("after_burst", j["after_burst"]))
            add(_fmt_stats("not_after_burst", j["not_after_burst"]))
            add("    (b) a burst printed while the trade was open")
            add(_fmt_stats("burst_during", j["burst_during"]))
            add(_fmt_stats("no_burst_during", j["no_burst_during"]))
            add("    (c) of (b), first burst vs the trade's direction")
            add(_fmt_stats("burst_agrees", j["burst_agrees"]))
            add(_fmt_stats("burst_opposes", j["burst_opposes"]))
            for grp in ("after_burst", "not_after_burst", "burst_during",
                        "no_burst_during", "burst_agrees", "burst_opposes"):
                join_rows.append({"k": k, "group": grp,
                                  "within_bars": args.within_bars,
                                  **j[grp]})

    text = "\n".join(lines)
    print(text)
    with open(os.path.join(args.out, "burst_map.txt"), "w") as fh:
        fh.write(text + "\n")
    pd.concat([v.assign(k=kk) for kk, v in tabs["burst_pct_by_k"].items()]
              ).to_csv(os.path.join(args.out, "hour_weekday_burst_pct.csv"))
    tabs["absmove"].to_csv(os.path.join(args.out, "hour_weekday_absmove.csv"))
    tabs["by_hour"].to_csv(os.path.join(args.out, "by_hour.csv"))
    tabs["by_weekday"].to_csv(os.path.join(args.out, "by_weekday.csv"))
    pd.concat([post_burst_drift(bars, k, bar_minutes=bm).assign(k=k)
               for k in ks]).to_csv(
        os.path.join(args.out, "post_burst_drift.csv"), index=False)
    moves.to_csv(os.path.join(args.out, "largest_moves.csv"), index=False)
    if join_rows:
        pd.DataFrame(join_rows).to_csv(
            os.path.join(args.out, "book_join.csv"), index=False)
    print(f"\n[written] {os.path.join(args.out, 'burst_map.txt')}")


if __name__ == "__main__":
    main()