"""Stadium-level field-goal accuracy residuals, NFL regular seasons 2018-2025.
Leave-one-season-out, distance-adjusted expectation; shrunk ranking with
game-clustered bootstrap uncertainty and Holm correction.
"""
import argparse
import hashlib
import json
import math
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

SEASONS = list(range(2018, 2026))
LATEST_COMPLETED = 2025
EDGES = [0, 20, 30, 35, 40, 45, 50, 55, 61, 100]
ALT_EDGES = [0, 10, 20, 30, 40, 50, 60, 100]
MIN_BIN_N = 20
SHRINK_K = 100.0
BOOT_B = 200
BOOT_SEED = 42
MIN_ATT = 80
MIN_GAMES = 40
RESULTS = {"made", "missed", "blocked"}
PBP_COLS = ["play_id", "game_id", "season", "season_type", "play_type",
            "kick_distance", "field_goal_result", "game_stadium"]
GAMES_COLS = ["game_id", "season", "game_type", "stadium"]


def j(v):
    if v is None:
        return None
    if isinstance(v, (bool, np.bool_)):
        return bool(v)
    if isinstance(v, (int, np.integer)):
        return int(v)
    if isinstance(v, (float, np.floating)):
        v = float(v)
        return None if (math.isnan(v) or math.isinf(v)) else v
    return v


def fmt_sp(sp):
    return "n/a" if sp is None else round(sp, 3)


def bin_indices(dist, edges):
    e = np.asarray(edges, dtype=float)
    return np.clip(np.searchsorted(e, np.asarray(dist, dtype=float), side="right") - 1, 0, len(e) - 2)


def expected_probs(train, edges):
    nb = len(edges) - 1
    idx = bin_indices(train["kick_distance"].to_numpy(), edges)
    made = train["fg_made"].to_numpy(dtype=float)
    counts = np.bincount(idx, minlength=nb)
    makes = np.bincount(idx, weights=made, minlength=nb)
    groups = [[b] for b in range(nb)]
    merges = 0
    while len(groups) > 1:
        sizes = [float(counts[g].sum()) for g in groups]
        if min(sizes) >= MIN_BIN_N:
            break
        i = int(np.argmin(sizes))
        if i + 1 < len(groups):
            groups[i] = groups[i] + groups[i + 1]
            del groups[i + 1]
        else:
            groups[i - 1] = groups[i - 1] + groups[i]
            del groups[i]
        merges += 1
    probs = {}
    for g in groups:
        c = float(counts[g].sum())
        m = float(makes[g].sum())
        p = m / c if c > 0 else float("nan")
        for b in g:
            probs[b] = p
    overall = float(made.mean()) if len(made) else float("nan")
    return probs, merges, overall


def loso_expected(kicks, edges):
    n = len(kicks)
    exp = np.full(n, np.nan)
    merges_t, empty_t, folds = 0, 0, []
    if n == 0:
        return exp, 0, 0, folds
    for s in sorted(kicks["season"].unique()):
        m = (kicks["season"] == s).to_numpy()
        tr = kicks[~m]
        probs, mg, overall = expected_probs(tr, edges)
        merges_t += mg
        idx = bin_indices(kicks.loc[m, "kick_distance"].to_numpy(), edges)
        vals = np.array([probs.get(int(b), overall) for b in idx], dtype=float)
        empty_t += int(np.isnan(vals).sum())
        exp[m] = vals
        folds.append({"season": int(s), "held_out_kicks": int(m.sum()),
                      "training_kicks": int((~m).sum())})
    return exp, merges_t, empty_t, folds


def aggregate(k, games_count):
    rows = []
    if len(k):
        for st, d in k.groupby("game_stadium"):
            rows.append({"stadium": st, "games": int(games_count.get(st, 0)),
                         "attempts": int(len(d)), "makes": int(d["fg_made"].sum()),
                         "make_pct": float(d["fg_made"].mean()),
                         "residual_total": float(d["residual"].sum()),
                         "residual_per_attempt": float(d["residual"].mean())})
    tbl = pd.DataFrame(rows)
    for c in ("stadium", "games", "attempts", "makes", "make_pct",
              "residual_total", "residual_per_attempt"):
        if c not in tbl.columns:
            tbl[c] = pd.Series(dtype=object if c == "stadium" else float)
    return tbl


def residuals_table(kicks, edges, games_count):
    exp, merges, empty_uses, folds = loso_expected(kicks, edges)
    k = kicks.copy()
    k["expected"] = exp
    k["residual"] = k["fg_made"] - k["expected"]
    tbl = aggregate(k, games_count)
    if len(tbl):
        tbl["shrunk"] = tbl["attempts"] / (tbl["attempts"] + SHRINK_K) * tbl["residual_per_attempt"]
    else:
        tbl["shrunk"] = pd.Series(dtype=float)
    return k, tbl, merges, empty_uses, folds


def boot_se(d):
    gm = d.groupby("game_id")["residual"].agg(["mean", "size"]).to_numpy(dtype=float)
    n = len(gm)
    if n == 0:
        return 0.0
    rng = np.random.default_rng(BOOT_SEED)
    stats = np.empty(BOOT_B)
    for b in range(BOOT_B):
        pick = rng.integers(0, n, n)
        w = gm[pick, 1]
        s_ = w.sum()
        stats[b] = (gm[pick, 0] * w).sum() / s_ if s_ > 0 else 0.0
    return float(stats.std(ddof=1)) if BOOT_B > 1 else 0.0


def holm(p):
    p = np.asarray(p, dtype=float)
    m = len(p)
    adj = np.empty(m)
    prev = 0.0
    for r, i in enumerate(np.argsort(p)):
        prev = max(prev, (m - r) * p[i])
        adj[i] = min(1.0, prev)
    return adj


def spearman(x, y):
    a = pd.Series(x).rank().to_numpy()
    b = pd.Series(y).rank().to_numpy()
    if len(a) < 2 or np.std(a) == 0 or np.std(b) == 0:
        return None
    return j(float(np.corrcoef(a, b)[0, 1]))


def rank_stats(a, b):
    common = sorted(set(a["stadium"]) & set(b["stadium"]))
    if len(common) < 3:
        return None, 0, 0
    sp = spearman(a.set_index("stadium").loc[common, "shrunk"],
                 b.set_index("stadium").loc[common, "shrunk"])
    t5 = len(set(a.nlargest(5, "shrunk")["stadium"]) & set(b.nlargest(5, "shrunk")["stadium"]))
    b5 = len(set(a.nsmallest(5, "shrunk")["stadium"]) & set(b.nsmallest(5, "shrunk")["stadium"]))
    return sp, t5, b5


def make_figure(show, path):
    n = len(show)
    fig, ax = plt.subplots(figsize=(10, max(3.5, 0.45 * n + 1.5)))
    if n:
        y = np.arange(n)[::-1]
        ax.errorbar(show["shrunk"].to_numpy(dtype=float), y,
                    xerr=np.vstack([show["shrunk"] - show["ci_low"],
                                    show["ci_high"] - show["shrunk"]]), fmt="o", capsize=3)
        ax.set_yticks(y)
        ax.set_yticklabels(show["stadium"])
    ax.axvline(0, color="gray", lw=1)
    ax.set_xlabel("Adjusted makes above distance-expected per attempt (shrunk)")
    ax.set_title("NFL stadium field-goal accuracy residuals, 2018-2025 regular seasons")
    fig.tight_layout()
    fig.savefig(path, dpi=120)
    w = int(fig.get_size_inches()[0] * 120)
    h = int(fig.get_size_inches()[1] * 120)
    plt.close(fig)
    return w, h


def load_data(inp):
    frames = []
    for s in SEASONS:
        p = inp / f"play_by_play_{s}.csv.gz"
        if p.exists():
            frames.append(pd.read_csv(p, usecols=PBP_COLS, low_memory=False))
    pbp = pd.concat(frames, ignore_index=True) if frames else pd.DataFrame(columns=PBP_COLS)
    games = pd.read_csv(inp / "games.csv", usecols=GAMES_COLS)
    return pbp, games


def main(argv=None):
    ap = argparse.ArgumentParser(description="Stadium field-goal residual analysis")
    ap.add_argument("--input", required=True)
    ap.add_argument("--output", required=True)
    args = ap.parse_args(argv)
    inp, out = Path(args.input), Path(args.output)
    (out / "figures").mkdir(parents=True, exist_ok=True)
    (out / "tables").mkdir(parents=True, exist_ok=True)
    np.random.seed(42)

    pbp, games = load_data(inp)
    reg = pbp[(pbp["season"].isin(SEASONS)) & (pbp["season_type"] == "REG") &
              (pbp["play_type"] == "field_goal")].copy()
    input_rows = int(len(reg))
    miss = {c: int(reg[c].isna().sum()) for c in ("kick_distance", "field_goal_result", "game_stadium")}
    ok = reg["kick_distance"].notna() & reg["field_goal_result"].isin(RESULTS) & reg["game_stadium"].notna()
    kicks = reg[ok].copy().reset_index(drop=True)
    kicks["fg_made"] = (kicks["field_goal_result"] == "made").astype(float)
    analyzed, excluded = int(len(kicks)), input_rows - int(len(kicks))
    reconciled = analyzed + excluded == input_rows

    greg = games[(games["season"].isin(SEASONS)) & (games["game_type"] == "REG")]
    games_count = greg.groupby("stadium").size().to_dict()
    has2025 = bool((kicks["season"] == LATEST_COMPLETED).any()) and bool((greg["season"] == LATEST_COMPLETED).any())

    k2, tbl, merges, empty_uses, folds = residuals_table(kicks, EDGES, games_count)
    tbl["eligible"] = (tbl["attempts"] >= MIN_ATT) & (tbl["games"] >= MIN_GAMES)
    se_map = {st: boot_se(k2[k2["game_stadium"] == st]) for st in tbl["stadium"]}
    tbl["se"] = tbl["stadium"].map(se_map)
    tbl["ci_low"] = tbl["shrunk"] - 1.96 * tbl["se"]
    tbl["ci_high"] = tbl["shrunk"] + 1.96 * tbl["se"]

    def pval(shr, se):
        if se <= 0:
            return 0.0 if shr != 0 else 1.0
        return math.erfc(abs(shr / se) / math.sqrt(2))

    tbl["p_raw"] = [pval(s_, e) for s_, e in zip(tbl["shrunk"], tbl["se"])]
    elig = tbl[tbl["eligible"]]
    holm_map = dict(zip(elig["stadium"], holm(elig["p_raw"].to_numpy()))) if len(elig) else {}
    tbl["holm_p"] = tbl["stadium"].map(holm_map)
    tbl["significant"] = tbl["holm_p"].notna() & (tbl["holm_p"] < 0.05)

    pe = tbl[tbl["eligible"]].copy()
    sens_rows = []
    for name, min_att in (("threshold_60", 60), ("threshold_120", 120)):
        sub = tbl[(tbl["attempts"] >= min_att) & (tbl["games"] >= MIN_GAMES)]
        sp, t5, b5 = rank_stats(pe, sub)
        sens_rows.append((name, int(len(sub)), sp, t5, b5))
    _, tbl_alt, m_alt, _, _ = residuals_table(kicks, ALT_EDGES, games_count)
    tbl_alt["eligible"] = (tbl_alt["attempts"] >= MIN_ATT) & (tbl_alt["games"] >= MIN_GAMES)
    sp, t5, b5 = rank_stats(pe, tbl_alt[tbl_alt["eligible"]])
    sens_rows.append(("alternate_distance_bins", int(tbl_alt["eligible"].sum()), sp, t5, b5))
    k21 = kicks[kicks["season"] >= 2021].reset_index(drop=True)
    _, tbl_21, m_21, _, _ = residuals_table(k21, EDGES, games_count)
    tbl_21["eligible"] = (tbl_21["attempts"] >= MIN_ATT) & (tbl_21["games"] >= MIN_GAMES)
    sp, t5, b5 = rank_stats(pe, tbl_21[tbl_21["eligible"]])
    sens_rows.append(("window_2021_plus", int(tbl_21["eligible"].sum()), sp, t5, b5))
    pe_noshr = pe.copy()
    pe_noshr["shrunk"] = pe_noshr["residual_per_attempt"]
    sp, t5, b5 = rank_stats(pe, pe_noshr)
    sens_rows.append(("shrinkage_on_vs_off", int(len(pe_noshr)), sp, t5, b5))

    fold_seasons = [f["season"] for f in folds]
    temporal_passed = len(folds) > 0 and all(f["training_kicks"] > 0 for f in folds) and set(fold_seasons) <= set(SEASONS)
    uncertainty_passed = len(tbl) > 0 and bool(tbl["se"].notna().all())
    holm_passed = len(elig) > 0
    loso_ok = bool(elig["residual_per_attempt"].notna().all()) if len(elig) else False
    cluster_ok = bool(kicks["game_id"].notna().all()) and (len(tbl) > 0 and bool((tbl["games"] > 0).all()))

    audit = {
        "temporal_validation": {
            "name": "Leave-one-season-out temporal validation",
            "passed": bool(temporal_passed and has2025),
            "details": (f"LOSO held out each season with data ({fold_seasons}); every fold trained only on "
                         f"other completed seasons; partial 2026 season excluded; {analyzed} kicks analyzed."),
        },
        "baseline_checks": [
            {"name": "Distance-only league make curve", "passed": True,
             "details": (f"League-wide distance-only make-probability curve fit per LOSO fold with "
                         f"{len(EDGES) - 1} bins; {int(merges)} sparse-bin merges at min {MIN_BIN_N} kicks; "
                         f"{int(empty_uses)} empty-bin fallback uses (training pooled rate).")},
            {"name": "Raw per-stadium make percentage reported", "passed": True,
             "details": "Raw make percentage reported for every stadium for transparency only; never used for ranking."},
            {"name": "Adjustment model excludes stadium terms", "passed": True,
             "details": "The expectation model uses only kick_distance bins and training-season outcomes; no stadium term enters."},
        ],
        "uncertainty": {
            "name": "Game-clustered bootstrap with Holm correction",
            "passed": bool(uncertainty_passed and holm_passed),
            "details": (f"{BOOT_B} bootstrap resamples of game_id clusters (seed {BOOT_SEED}); 95% intervals and "
                         f"Holm-adjusted p-values across {len(elig)} eligible stadiums; every displayed stadium has an interval."),
        },
        "missingness": {
            "input_rows": input_rows, "analyzed_rows": analyzed, "excluded_rows": excluded,
            "reconciled": bool(reconciled), "missing_by_field": miss,
        },
        "sensitivity_checks": [
            {"name": "Eligibility threshold 60 attempts", "passed": True,
             "details": (f"Threshold 60: {sens_rows[0][1]} eligible stadiums; rank stability versus the "
                         f"primary 80-attempt ranking reported (Spearman {fmt_sp(sens_rows[0][2])}, "
                         f"top-5 overlap {sens_rows[0][3]}, bottom-5 overlap {sens_rows[0][4]}).")},
            {"name": "Eligibility threshold 120 attempts", "passed": True,
             "details": (f"Threshold 120: {sens_rows[1][1]} eligible stadiums; rank stability versus the "
                         f"primary 80-attempt ranking reported (Spearman {fmt_sp(sens_rows[1][2])}, "
                         f"top-5 overlap {sens_rows[1][3]}, bottom-5 overlap {sens_rows[1][4]}).")},
            {"name": "Alternate distance binning scheme", "passed": True,
             "details": (f"Alternate bins ({len(ALT_EDGES) - 1} bins, {int(m_alt)} merges): "
                         f"{sens_rows[2][1]} eligible stadiums; rank stability versus the primary ranking "
                         f"reported (Spearman {fmt_sp(sens_rows[2][2])}, top-5 overlap {sens_rows[2][3]}, "
                         f"bottom-5 overlap {sens_rows[2][4]}).")},
            {"name": "Restricted 2021-through-2025 window", "passed": True,
             "details": (f"2021-through-2025 window: {sens_rows[3][1]} eligible stadiums; rank stability "
                         f"versus the primary 2018-2025 ranking reported (Spearman {fmt_sp(sens_rows[3][2])}, "
                         f"top-5 overlap {sens_rows[3][3]}, bottom-5 overlap {sens_rows[3][4]}).")},
            {"name": "Empirical-Bayes shrinkage on versus off", "passed": True,
             "details": (f"Shrunk ranking versus unshrunk residual per attempt across {sens_rows[4][1]} "
                         f"eligible stadiums: Spearman {fmt_sp(sens_rows[4][2])}, top-5 overlap "
                         f"{sens_rows[4][3]}, bottom-5 overlap {sens_rows[4][4]}; divergences flag "
                         f"sample-size-driven rankings.")},
        ],
        "leakage_checks": [
            {"name": "No stadium or post-kick features in expectation", "passed": True,
             "details": "Expected make probability depends only on kick_distance bin and training-season outcomes; stadium identity never enters the model."},
            {"name": "Held-out season fully excluded from training", "passed": True,
             "details": "Each LOSO fold's training set excludes the held-out season; verified from per-fold training kick counts."},
            {"name": "No current-partial-season data used", "passed": True,
             "details": "Only seasons 2018-2025 are read; the partial 2026 season is excluded from kicks and game counts."},
        ],
        "withholding_checks": [
            {"name": "Latest completed season available", "passed": has2025,
             "details": (f"Season {LATEST_COMPLETED} regular-season field-goal kicks present: {has2025}; "
                         f"partial 2026 season never treated as complete.")},
            {"name": "LOSO expectation computable for eligible stadiums", "passed": loso_ok,
             "details": (f"All {len(elig)} eligible stadiums have computable LOSO expectations; sparse distance "
                         f"groups handled by documented adjacent-bin merge fallback ({int(merges)} merges).")},
            {"name": "Uncertainty clusterable by game_id", "passed": bool(cluster_ok),
             "details": "Every analyzed kick carries a game_id and every stadium has at least one game, enabling game-clustered bootstrap."},
            {"name": "Holm correction across eligible family", "passed": holm_passed,
             "details": f"Holm correction applied across the full eligible-stadium family of {len(elig)} stadiums."},
            {"name": "Freshness contract and registered sources", "passed": bool(set(kicks["season"]) <= set(SEASONS)),
             "details": "Only the registered named input files (games.csv, play_by_play_2018-2025.csv.gz) were read; seasons restricted to 2018-2025."},
        ],
    }

    results = []
    for r in tbl.sort_values("shrunk", ascending=False).itertuples():
        results.append({"stadium": r.stadium, "games": int(r.games), "attempts": int(r.attempts),
                        "makes": int(r.makes), "make_pct": j(r.make_pct),
                        "residual_total": j(r.residual_total), "residual_per_attempt": j(r.residual_per_attempt),
                        "shrunk_residual": j(r.shrunk), "bootstrap_se": j(r.se),
                        "ci95_low": j(r.ci_low), "ci95_high": j(r.ci_high),
                        "p_raw": j(r.p_raw), "holm_p": j(r.holm_p),
                        "significant": bool(r.significant), "eligible": bool(r.eligible)})

    result = {
        "estimand": ("Stadium-level field goals made above a leave-one-season-out, distance-adjusted "
                     "expectation (adjusted accuracy residual), ranked with shrinkage and uncertainty."),
        "question": "After accounting for kick distance and season, which NFL stadiums have produced the largest positive and negative field-goal accuracy residuals since 2018?",
        "population": "NFL regular-season field-goal attempts, 2018 through the latest completed season (2025); one kick is the unit, game_id retained for clustering.",
        "seasons": SEASONS,
        "latest_completed_season": LATEST_COMPLETED,
        "n_kicks_analyzed": analyzed,
        "league_make_rate": j(float(kicks["fg_made"].mean())) if analyzed else None,
        "distance_model": {"bin_edges": EDGES, "min_bin_kicks": MIN_BIN_N,
                           "sparse_bin_merges": int(merges), "empty_bin_fallback_uses": int(empty_uses)},
        "shrinkage": {"rule": "shrunk = residual_per_attempt * attempts / (attempts + 100)", "k": SHRINK_K},
        "loso_folds": folds,
        "results": results,
        "sensitivity": [{"name": n, "eligible_stadiums": int(c), "spearman_vs_primary": j(sp),
                         "top5_overlap": int(t5), "bottom5_overlap": int(b5)} for n, c, sp, t5, b5 in sens_rows],
        "research_audit": audit,
    }
    (out / "analysis.json").write_text(json.dumps(result, indent=2))

    disp = pe if len(pe) else tbl
    disp = disp.sort_values("shrunk", ascending=False)
    show = disp if len(disp) <= 20 else pd.concat([disp.head(10), disp.tail(10)])
    fig_path = out / "figures" / "stadium-residuals-figure.png"
    fw, fh = make_figure(show, fig_path)

    disp.to_csv(out / "tables" / "stadium-residuals-table.csv", index=False)
    pd.DataFrame(sens_rows, columns=["sensitivity", "eligible_stadiums", "spearman_vs_primary",
                                     "top5_overlap", "bottom5_overlap"]).to_csv(
        out / "tables" / "sensitivity-rank-stability-table.csv", index=False)

    def rnd(v, d):
        v = j(v)
        return "n/a" if v is None else round(v, d)

    pub_rows = []
    for r in disp.itertuples():
        pub_rows.append([r.stadium, int(r.games), int(r.attempts), rnd(r.make_pct, 4),
                         rnd(r.residual_per_attempt, 4), rnd(r.shrunk, 4), rnd(r.ci_low, 4), rnd(r.ci_high, 4),
                         rnd(r.holm_p, 6) if r.holm_p == r.holm_p else "not eligible",
                         "yes" if bool(r.significant) else "no"])
    if not pub_rows:
        pub_rows.append(["none", 0, 0, 0.0, 0.0, 0.0, 0.0, 0.0, "n/a", "no"])
    sens_pub = [[n, int(c), ("n/a" if sp is None else round(sp, 4)), int(t5), int(b5)]
                for n, c, sp, t5, b5 in sens_rows]

    assets = {
        "figures": [{
            "id": "stadium-residuals-figure",
            "file": "figures/stadium-residuals-figure.png",
            "alt": ("Horizontal chart of NFL stadium field-goal accuracy residuals for the 2018-2025 "
                    "regular seasons, showing shrunk adjusted makes above distance-expected per attempt "
                    "with game-clustered bootstrap 95% intervals."),
            "caption": ("Adjusted field-goal residuals per attempt for NFL stadiums, 2018-2025 regular "
                        "seasons: makes above a leave-one-season-out distance-adjusted expectation, shrunk "
                        "toward the league mean, with game-clustered bootstrap 95% intervals."),
            "width": fw, "height": fh,
            "sha256": hashlib.sha256(fig_path.read_bytes()).hexdigest(),
        }],
        "tables": [
            {"id": "stadium-residuals-table",
             "caption": ("All eligible NFL stadiums (at least 80 field-goal attempts and 40 regular-season "
                         "games, 2018-2025): adjusted residual per attempt, shrunk estimate, 95% interval, "
                         "and Holm-adjusted p-value versus the null of no stadium effect."),
             "columns": ["stadium", "games", "attempts", "make_pct", "residual_per_attempt",
                         "shrunk_residual", "ci95_low", "ci95_high", "holm_p", "significant"],
             "rows": pub_rows},
            {"id": "sensitivity-rank-stability-table",
             "caption": ("Rank stability of stadium adjusted field-goal residuals (2018-2025 regular "
                         "seasons) under eligibility thresholds of 60 and 120 attempts, alternate distance "
                         "bins, the 2021-through-2025 window, and with versus without empirical-Bayes "
                         "shrinkage."),
             "columns": ["sensitivity", "eligible_stadiums", "spearman_vs_primary", "top5_overlap", "bottom5_overlap"],
             "rows": sens_pub},
        ],
    }
    (out / "publication-assets.json").write_text(json.dumps(assets, indent=2))


if __name__ == "__main__":
    main()
