"""Robustness battery (Table VIII) from stored event files.

For each subsample/design variant, re-estimate (i) the market-wide baseline
ΔΔ breadth at h ∈ {5, 30} and (ii) the headline network differential at h=5,
with day-block bootstrap SEs:

  full                 full sample (reference)
  post-decimal         events dated 2002+ (kills 1/8–1/16 tick regime)
  mid-day              event minute in [10:00, 15:30) ET
  ex-announcement      exclude 10:00±10min and 14:00–14:30 ET event minutes
  calm-years           exclude 2008, 2009, 2020
  crisis-years         2008, 2009, 2020 only
  |z|>=8, |z|>=12      nested source-magnitude thresholds
  ex-noncommon         exclude sources flagged non-common (ticker heuristic)
  placebo-network      network differential computed on placebo events only
                       (falsification: should be ≈ 0)
  month-cluster        full sample, bootstrap clustered by month instead of day

Outputs: outputs/estimates/stage_f_robustness.json,
         outputs/tables/table8_robustness.{tex,csv}
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import polars as pl

ROOT = Path(__file__).resolve().parents[1]
EV = ROOT / "outputs" / "estimates" / "events"
EV2 = ROOT / "outputs" / "estimates" / "events2"
UNIV = ROOT / "data" / "metadata" / "universe_monthly.parquet"
TAB = ROOT / "outputs" / "tables"
B = 500
SEED = 20260713


def dd_boot(df, col, rng, cluster="date"):
    g = (df.filter(pl.col(col).is_not_null())
         .group_by(cluster, "kind")
         .agg(pl.col(col).sum().alias("s"), pl.len().alias("n"))
         .pivot(on="kind", index=cluster, values=["s", "n"]))
    need = ("s_real", "n_real", "s_placebo", "n_placebo")
    if any(c not in g.columns for c in need) or g.height < 10:
        return np.nan, np.nan, 0
    g = g.fill_null(0)
    sr, nr = g["s_real"].to_numpy(), g["n_real"].to_numpy()
    sp, npb = g["s_placebo"].to_numpy(), g["n_placebo"].to_numpy()
    est = sr.sum() / max(nr.sum(), 1) - sp.sum() / max(npb.sum(), 1)
    idx = rng.integers(0, len(sr), size=(B, len(sr)))
    bs = (sr[idx].sum(1) / np.maximum(nr[idx].sum(1), 1)
          - sp[idx].sum(1) / np.maximum(npb[idx].sum(1), 1))
    return float(est), float(bs.std()), int(nr.sum())


def one_boot(df, col, rng, kind="placebo"):
    """Single-arm mean with day bootstrap (for the placebo-network check)."""
    g = (df.filter((pl.col("kind") == kind) & pl.col(col).is_not_null())
         .group_by("date").agg(pl.col(col).sum().alias("s"), pl.len().alias("n")))
    s, n = g["s"].to_numpy(), g["n"].to_numpy()
    est = s.sum() / max(n.sum(), 1)
    idx = rng.integers(0, len(s), size=(B, len(s)))
    bs = s[idx].sum(1) / np.maximum(n[idx].sum(1), 1)
    return float(est), float(bs.std()), int(n.sum())


def main() -> None:
    rng = np.random.default_rng(SEED)
    e1 = pl.concat([pl.read_parquet(p) for p in sorted(EV.glob("events_*.parquet"))
                    if p.stat().st_size > 500], how="diagonal").with_columns(
        pl.col("date").str.slice(0, 4).cast(pl.Int32).alias("year"),
        pl.col("z").abs().alias("absz"),
        (pl.col("breadth_5") - pl.col("pre_breadth")).alias("db_5"),
        (pl.col("breadth_30") - pl.col("pre_breadth")).alias("db_30"),
        pl.col("date").str.slice(0, 7).alias("ym"))
    e2 = pl.concat([pl.read_parquet(p) for p in sorted(EV2.glob("events2_*.parquet"))
                    if p.stat().st_size > 500], how="diagonal").with_columns(
        pl.col("date").str.slice(0, 4).cast(pl.Int32).alias("year"),
        pl.col("z").abs().alias("absz"),
        (pl.col("nbh_5") - pl.col("non_5")).alias("nd_5"),
        pl.col("date").str.slice(0, 7).alias("ym"))
    noncommon = (pl.read_parquet(UNIV)
                 .filter(pl.col("flag_noncommon"))
                 .select("month", "ticker").unique()
                 .with_columns(pl.lit(True).alias("nc")))
    e1 = e1.join(noncommon, on=["month", "ticker"], how="left")
    e2 = e2.join(noncommon, on=["month", "ticker"], how="left")

    CRISIS = [2008, 2009, 2020]
    subs = {
        "Full sample": (pl.lit(True), pl.lit(True)),
        "Post-decimalization (2002+)": (pl.col("year") >= 2002,) * 2,
        "Mid-day events (10:00–15:30)": (pl.col("mod").is_between(600, 929),) * 2,
        "Ex announcement windows": (
            ~(pl.col("mod").is_between(595, 615)
              | pl.col("mod").is_between(840, 870)),) * 2,
        "Calm years (ex 2008–09, 2020)": (~pl.col("year").is_in(CRISIS),) * 2,
        "Crisis years only": (pl.col("year").is_in(CRISIS),) * 2,
        "Source $|z|\\ge 8$": (pl.col("absz") >= 8,) * 2,
        "Source $|z|\\ge 12$": (pl.col("absz") >= 12,) * 2,
        "Ex non-common tickers": (pl.col("nc").is_null(),) * 2,
    }
    rows = []
    for label, (c1, c2) in subs.items():
        b5, s5, n5 = dd_boot(e1.filter(c1), "db_5", rng)
        b30, s30, _ = dd_boot(e1.filter(c1), "db_30", rng)
        nd, snd, nn = dd_boot(e2.filter(c2), "nd_5", rng)
        rows.append({"sample": label, "db5": b5, "db5_se": s5,
                     "db30": b30, "db30_se": s30,
                     "nd5": nd, "nd5_se": snd, "n": n5})
        print(f"{label:<32} ΔΔb5={b5*1e4:+.2f}({s5*1e4:.2f}) "
              f"ΔΔb30={b30*1e4:+.2f}({s30*1e4:.2f}) "
              f"net5={nd*1e4:+.1f}({snd*1e4:.1f}) n={n5:,}", flush=True)

    # month-cluster bootstrap variant (full sample)
    b5, s5, n5 = dd_boot(e1, "db_5", rng, cluster="ym")
    nd, snd, _ = dd_boot(e2, "nd_5", rng, cluster="ym")
    rows.append({"sample": "Month-cluster bootstrap", "db5": b5, "db5_se": s5,
                 "db30": None, "db30_se": None, "nd5": nd, "nd5_se": snd, "n": n5})
    print(f"month-cluster: ΔΔb5={b5*1e4:+.2f}({s5*1e4:.2f}) net5={nd*1e4:+.1f}({snd*1e4:.1f})")

    # placebo-network falsification: nbh−non among placebos alone
    pe, ps, pn = one_boot(e2, "nd_5", rng, kind="placebo")
    rows.append({"sample": "Placebo events, neighbor diff", "db5": None,
                 "db5_se": None, "db30": None, "db30_se": None,
                 "nd5": pe, "nd5_se": ps, "n": pn})
    print(f"placebo network diff: {pe*1e4:+.1f}({ps*1e4:.1f})")

    out = pl.DataFrame(rows)
    out.write_csv(TAB / "table8_robustness.csv")
    (ROOT / "outputs" / "estimates" / "stage_f_robustness.json").write_text(
        json.dumps(rows, indent=1))

    def cell(v, s):
        if v is None:
            return "---"
        return f"{v*1e4:,.2f} ({s*1e4:,.2f})"
    lines = [
        "% Auto-generated by scripts/16_robustness.py — do not edit numbers.",
        "\\begin{table}[htbp]\\centering\\begin{adjustbox}{max width=\\textwidth}\\begin{threeparttable}",
        "\\caption{Robustness of the market-wide null and the network channel}",
        "\\label{tab:robustness}",
        "\\small\\begin{tabular}{lrrr}",
        "\\toprule",
        "Sample / design & $\\Delta\\Delta$Breadth$_{5}$ & "
        "$\\Delta\\Delta$Breadth$_{30}$ & Network diff$_{5}$ \\\\",
        "\\midrule"]
    for r in rows:
        lines.append(f"{r['sample']} & {cell(r['db5'], r['db5_se'])} & "
                     f"{cell(r['db30'], r['db30_se'])} & "
                     f"{cell(r['nd5'], r['nd5_se'])} \\\\")
    lines += [
        "\\bottomrule\\end{tabular}",
        "\\begin{tablenotes}[flushleft]\\footnotesize",
        "\\item Notes: All entries $\\times 10^{-4}$, day-block bootstrap SEs in",
        "parentheses (500 reps; final row of the market-wide block uses",
        "month-cluster resampling). $\\Delta\\Delta$Breadth = abnormal market-wide",
        "breadth (real $-$ placebo, post $-$ pre). Network diff = neighbor $-$",
        "non-neighbor abnormal breadth at $h{=}5$. `Placebo events' row reports",
        "the neighbor differential among placebo events alone (falsification;",
        "expected $\\approx 0$).",
        "\\end{tablenotes}\\end{threeparttable}\\end{adjustbox}\\end{table}"]
    (TAB / "table8_robustness.tex").write_text("\n".join(lines))
    (ROOT / "paper" / "tables" / "table8_robustness.tex").write_text("\n".join(lines))
    print("robustness table written")


if __name__ == "__main__":
    main()