"""
Tutorial 93 - Multiple comparisons: what survives a 30-test screen.

One p-value below 0.05 is a finding. Thirty p-values below 0.05, computed at
the same time on the same season, are a lottery ticket - and the fix is not a
matter of taste, it is arithmetic. This tutorial builds the exact two-sided
binomial test by hand, runs it once, then runs it thirty times across the
bundled 2023-24 NBA season and asks the only question that matters: how many
of those thirty "findings" would a league with nothing going on have produced
anyway? It then applies the two standard corrections - Bonferroni for the
family-wise error rate and Benjamini-Hochberg (1995) for the false discovery
rate - and turns the same machinery on two further screens, one of 30 tests
and one of 191, where the answer comes out the other way.

Runs entirely offline from the two bundled CSVs next to this script (no API,
no key, no scipy). Run:

    python downloads/93_multiple_comparisons_bonferroni_and_fdr.py

Data: Basketball-Reference 2023-24 schedule (nba_home_results.csv) and the
public NBA_Shots_04_25 shot log, 25,000-shot sample (nba_league_shots.csv).
"""
import math
import os

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

import sdt_common as sdt

sdt.init("multiple-comparisons-bonferroni-and-fdr")
HERE = os.path.dirname(os.path.abspath(__file__))
GAMES_CSV = os.path.join(HERE, "nba_home_results.csv")
SHOTS_CSV = os.path.join(HERE, "nba_league_shots.csv")

IST_FINAL = "2023-12-09"   # the In-Season Tournament final, played in Las Vegas


# --- the exact tests, written out in full ------------------------------------------
def log_pmf(k, n, p):
    """log of the binomial probability of exactly k successes in n trials."""
    return (math.lgamma(n + 1) - math.lgamma(k + 1) - math.lgamma(n - k + 1)
            + k * math.log(p) + (n - k) * math.log1p(-p))


def binom_test(k, n, p):
    """Exact two-sided binomial p-value: the total probability of any outcome at
    least as unlikely as the one observed (the 'method of small p-values')."""
    obs = log_pmf(k, n, p) + 1e-9
    return min(1.0, sum(math.exp(log_pmf(i, n, p))
                        for i in range(n + 1) if log_pmf(i, n, p) <= obs))


def fisher_exact(a, b, c, d):
    """Exact two-sided p-value for the 2x2 table [[a, b], [c, d]]."""
    n, r1, r2, c1 = a + b + c + d, a + b, c + d, a + c

    def prob(x):
        return math.comb(r1, x) * math.comb(r2, c1 - x) / math.comb(n, c1)

    obs = prob(a) * (1 + 1e-9)
    return min(1.0, sum(prob(x) for x in range(max(0, c1 - r2), min(r1, c1) + 1)
                        if prob(x) <= obs))


def bh_survivors(pvals, q):
    """Benjamini-Hochberg: how many hypotheses are rejected at FDR level q."""
    p = np.sort(np.asarray(pvals, dtype=float))
    m = len(p)
    below = np.nonzero(p <= (np.arange(1, m + 1) / m) * q)[0]
    return int(below.max() + 1) if len(below) else 0


# --- 1. one test, done properly ----------------------------------------------------
games = pd.read_csv(GAMES_CSV)
season = games[games.date != IST_FINAL].copy()      # neutral-site game, not a home game
season["home_win"] = (season.home_pts > season.away_pts).astype(int)
p_home = season.home_win.mean()

bos = season[season.home_team == "Boston Celtics"]
k_bos, n_bos = int(bos.home_win.sum()), len(bos)
p_bos = binom_test(k_bos, n_bos, p_home)
extreme = sum(1 for i in range(n_bos + 1)
              if log_pmf(i, n_bos, p_home) <= log_pmf(k_bos, n_bos, p_home) + 1e-9)

with sdt.snippet("load"):
    print(f"{len(games)} rows; dropping the {IST_FINAL} neutral-site final leaves "
          f"{len(season)} home games")
    print(f"home teams won {int(season.home_win.sum())} of {len(season)} = "
          f"{100 * p_home:.1f}%   <- the benchmark every team is tested against\n")
    print(f"Boston went {k_bos}-{n_bos - k_bos} at home ({100 * k_bos / n_bos:.1f}%)")
    print(f"  P(exactly {k_bos} of {n_bos} at {p_home:.4f}) = "
          f"{math.exp(log_pmf(k_bos, n_bos, p_home)):.3e}")
    print(f"  outcomes at least that unlikely: {extreme} of the {n_bos + 1} possible totals")
    print(f"  exact two-sided p = {p_bos:.2e}")

assert len(games) == 1231 and len(season) == 1230
assert int(season.home_win.sum()) == 668
assert abs(p_home - 668 / 1230) < 1e-12
assert (k_bos, n_bos) == (37, 41)
assert 9.0e-7 < p_bos < 9.9e-7, p_bos
assert extreme == 12
assert binom_test(int(season.home_win.sum()), len(season), 0.5) < 0.003
assert abs(binom_test(20, 40, 0.5) - 1.0) < 1e-9        # the dead-centre outcome
assert abs(binom_test(0, 10, 0.5) - 2 * 0.5 ** 10) < 1e-12
assert binom_test(30, 41, p_home) > binom_test(37, 41, p_home)

# --- 2. thirty tests at once -------------------------------------------------------
rows = []
for team, sub in season.groupby("home_team"):
    n, k = len(sub), int(sub.home_win.sum())
    rows.append((team, n, k, 100 * k / n, binom_test(k, n, p_home)))
screen = pd.DataFrame(rows, columns=["team", "games", "wins", "win_pct", "p"])
screen = screen.sort_values("p").reset_index(drop=True)
m = len(screen)
raw_hits = int((screen.p < 0.05).sum())

with sdt.snippet("screen"):
    print(f"the same test, run {m} times - once per home team\n")
    top = screen.head(10).copy()
    top["win_pct"] = top.win_pct.map("{:.1f}".format)
    top["p"] = top.p.map("{:.2e}".format)
    print(top.to_string(index=False))
    print(f"...\n{raw_hits} of {m} teams come in under p < 0.05")
    print(f"tests expected under 0.05 by chance alone: {m * 0.05:.1f}")

assert m == 30 and raw_hits == 13
assert int(screen.games.sum()) == 1230 and int(screen.wins.sum()) == 668
assert screen.team.iloc[0] == "Boston Celtics"
assert screen.team.iloc[1] == "Washington Wizards"
assert (screen.p.diff().dropna() >= 0).all()             # sorted ascending
assert int((screen.p < 0.01).sum()) == 9
assert 0.99 < screen.p.iloc[-1] <= 1.0                   # Miami sits on the league rate
assert screen.games.between(40, 42).all()

# --- 3. what a league with nothing going on produces --------------------------------
rng = np.random.default_rng(93)
sizes = screen.games.to_numpy()
lut = {int(n): np.array([binom_test(k, int(n), p_home) for k in range(int(n) + 1)])
       for n in np.unique(sizes)}
REPS = 20000
draws = rng.binomial(sizes, p_home, size=(REPS, m))
sim_p = np.empty(draws.shape, dtype=float)
for j, n in enumerate(sizes):
    sim_p[:, j] = lut[int(n)][draws[:, j]]
sim_hits = (sim_p < 0.05).sum(axis=1)
fwer_sim = float((sim_hits >= 1).mean())
fwer_naive = 1 - 0.95 ** m

with sdt.snippet("nullsim"):
    print(f"{REPS} simulated seasons in which all {m} teams are identical")
    print(f"(every home game an independent {100 * p_home:.1f}% coin)\n")
    print(f"average teams flagged at p < 0.05    : {sim_hits.mean():.2f}")
    print(f"seasons flagging at least one team   : {100 * fwer_sim:.1f}%")
    print(f"    what 1 - 0.95^{m} predicts        : {100 * fwer_naive:.1f}%")
    print(f"seasons flagging three or more       : {100 * (sim_hits >= 3).mean():.1f}%")
    print(f"seasons flagging {raw_hits} or more (as we saw) : "
          f"{int((sim_hits >= raw_hits).sum())} of {REPS}")
    print(f"busiest simulated season                : {int(sim_hits.max())} flags")
    print("\nflags per season:  " + "   ".join(
        f"{c}: {int((sim_hits == c).sum())}" for c in range(6)))

assert 1.15 < sim_hits.mean() < 1.30, sim_hits.mean()
assert 0.69 < fwer_sim < 0.73, fwer_sim
assert fwer_sim < fwer_naive                    # exact discrete tests are conservative
assert 0.78 < fwer_naive < 0.79
assert int((sim_hits >= raw_hits).sum()) == 0
assert int(sim_hits.max()) < raw_hits
assert 0.10 < (sim_hits >= 3).mean() < 0.14
assert int((sim_hits == 0).sum()) > 5000

# --- 4. the two corrections --------------------------------------------------------
ALPHA = 0.05
bonf_cut = ALPHA / m
bonf_hits = int((screen.p < bonf_cut).sum())
bh_hits = bh_survivors(screen.p, ALPHA)
ladder = screen.head(16).copy()
ladder["i"] = np.arange(1, len(ladder) + 1)
ladder["bh_crit"] = ladder["i"] / m * ALPHA
ladder["passes"] = np.where(ladder.p <= ladder.bh_crit, "yes", "no")

with sdt.snippet("correct"):
    print(f"m = {m} tests at alpha = {ALPHA}\n")
    print(f"Bonferroni: reject when p < alpha/m = {bonf_cut:.6f}  ->  {bonf_hits} survive")
    print(f"Benjamini-Hochberg: reject when p <= (i/m) x q       ->  {bh_hits} survive\n")
    show = ladder[["i", "team", "p", "bh_crit", "passes"]].copy()
    show["p"] = show.p.map("{:.2e}".format)
    show["bh_crit"] = show.bh_crit.map("{:.5f}".format)
    print(show.to_string(index=False))
    print(f"\nlargest i whose p clears its own step: {bh_hits} -> reject ranks 1..{bh_hits}")
    print(f"uncorrected {raw_hits}   BH(q=0.05) {bh_hits}   "
          f"BH(q=0.10) {bh_survivors(screen.p, 0.10)}   Bonferroni {bonf_hits}")

assert abs(bonf_cut - 0.05 / 30) < 1e-12
assert bonf_hits == 9 and bh_hits == 12
assert bh_survivors(screen.p, 0.10) == 13
assert bh_hits >= bonf_hits                      # BH is never stricter than Bonferroni
assert bh_hits <= raw_hits                       # and never looser than no correction
assert bh_survivors([0.9, 0.8, 0.7], 0.05) == 0
assert bh_survivors([0.001] * 10, 0.05) == 10
assert abs(float(ladder.bh_crit.iloc[11]) - 12 / 30 * 0.05) < 1e-12
assert float(screen.p.iloc[11]) <= 12 / 30 * 0.05
assert float(screen.p.iloc[12]) > 13 / 30 * 0.05
assert screen.team.iloc[12] == "Orlando Magic"   # flagged raw, dropped by BH

# --- 5. the same machinery on the question that actually matters --------------------
hca_rows = []
for team in sorted(season.home_team.unique()):
    h = season[season.home_team == team]
    a = season[season.away_team == team]
    hw = int(h.home_win.sum())
    aw = int((a.away_pts > a.home_pts).sum())
    hca_rows.append((team, hw, len(h) - hw, aw, len(a) - aw,
                     100 * (hw / len(h) - aw / len(a)),
                     fisher_exact(hw, len(h) - hw, aw, len(a) - aw)))
hca = pd.DataFrame(hca_rows, columns=["team", "hw", "hl", "aw", "al", "gap_pp", "p"])
hca = hca.sort_values("p").reset_index(drop=True)
hca_raw = int((hca.p < 0.05).sum())
hca_bonf = int((hca.p < ALPHA / m).sum())
hca_bh = bh_survivors(hca.p, ALPHA)
pooled = binom_test(int(season.home_win.sum()), len(season), 0.5)

rng2 = np.random.default_rng(931)
nh = np.array([len(season[season.home_team == t]) for t in sorted(season.home_team.unique())])
na = np.array([len(season[season.away_team == t]) for t in sorted(season.home_team.unique())])
REPS2 = 4000
cache = {}


def fisher_cached(a, b, c, d):
    key = (a, b, c, d)
    if key not in cache:
        cache[key] = fisher_exact(a, b, c, d)
    return cache[key]


hw_sim = rng2.binomial(nh, p_home, size=(REPS2, m))
aw_sim = rng2.binomial(na, 1 - p_home, size=(REPS2, m))
hits2 = np.zeros(REPS2, dtype=int)
for r in range(REPS2):
    hits2[r] = sum(
        fisher_cached(int(hw_sim[r, j]), int(nh[j] - hw_sim[r, j]),
                      int(aw_sim[r, j]), int(na[j] - aw_sim[r, j])) < 0.05
        for j in range(m))

with sdt.snippet("hca"):
    print("screen 2: does THIS team have a home-court advantage of its own?")
    print("(exact 2x2 test of each team's home record against its own road record)\n")
    top = hca.head(8).copy()
    top["gap_pp"] = top.gap_pp.map("{:+.1f}".format)
    top["p"] = top.p.map("{:.4f}".format)
    print(top.to_string(index=False))
    print(f"...\nteams with a positive home-road gap: {int((hca.gap_pp > 0).sum())} of {m}")
    print(f"uncorrected p < 0.05: {hca_raw}    Bonferroni: {hca_bonf}    "
          f"BH(q=0.05): {hca_bh}    BH(q=0.10): {bh_survivors(hca.p, 0.10)}")
    print(f"\nleague-wide, pooled over all {len(season)} games: "
          f"{int(season.home_win.sum())}-{len(season) - int(season.home_win.sum())}, "
          f"exact p = {pooled:.4f}")
    print(f"{REPS2} simulated leagues where every team shares the SAME home edge:")
    print(f"  teams flagged, average {hits2.mean():.2f}; "
          f"{hca_raw} or more happened {100 * (hits2 >= hca_raw).mean():.1f}% of the time")

assert len(hca) == m
assert hca_raw == 5 and hca_bonf == 0 and hca_bh == 0
assert bh_survivors(hca.p, 0.10) == 0
assert bh_survivors(hca.p, 0.20) == 5
assert int((hca.gap_pp > 0).sum()) == 23 and int((hca.gap_pp < 0).sum()) == 6
assert hca.team.iloc[0] == "Houston Rockets"
assert 0.0075 < float(hca.p.iloc[0]) < 0.0078
assert 0.0025 < pooled < 0.0030, pooled
assert pooled < 0.01                             # one pre-specified test, nothing to correct
assert pooled < float(hca.p.iloc[0])             # the pooled test beats every team's own
assert 2.5 < hits2.mean() < 3.3, hits2.mean()
assert 0.10 < (hits2 >= hca_raw).mean() < 0.25
assert abs(fisher_exact(10, 10, 10, 10) - 1.0) < 1e-9
assert fisher_exact(20, 0, 0, 20) < 1e-9

# --- 6. scale it up: hundreds of tests ---------------------------------------------
shots = pd.read_csv(SHOTS_CSV)
lg_fg = shots.SHOT_MADE.mean()
fg = shots.groupby("PLAYER_NAME").SHOT_MADE.agg(att="size", made="sum")
fg = fg[fg.att >= 50].copy()
fg["pct"] = 100 * fg.made / fg.att
fg["p"] = [binom_test(int(k), int(n), lg_fg) for n, k in zip(fg.att, fg.made)]
fg = fg.sort_values("p")
m_fg = len(fg)

threes = shots[shots.SHOT_TYPE == "3PT Field Goal"]
lg_3p = threes.SHOT_MADE.mean()
tp = threes.groupby("PLAYER_NAME").SHOT_MADE.agg(att="size", made="sum")
tp = tp[tp.att >= 30].copy()
tp["pct"] = 100 * tp.made / tp.att
tp["p"] = [binom_test(int(k), int(n), lg_3p) for n, k in zip(tp.att, tp.made)]
tp = tp.sort_values("p")
m_tp = len(tp)

survivors_fg = list(fg.index[:bh_survivors(fg.p, 0.05)])
rim = shots.assign(rim=(shots.BASIC_ZONE == "Restricted Area").astype(int))
rim_share = rim.groupby("PLAYER_NAME").rim.mean()
rim_share_lg = rim.rim.mean()

with sdt.snippet("shooters"):
    print(f"{len(shots)} shots, {shots.PLAYER_NAME.nunique()} shooters; "
          f"league {100 * lg_fg:.1f}% from the field, {100 * lg_3p:.1f}% from three\n")
    for label, tbl, mm, base in [("all field goals, 50+ attempts", fg, m_fg, lg_fg),
                                 ("threes only, 30+ attempts", tp, m_tp, lg_3p)]:
        raw = int((tbl.p < 0.05).sum())
        print(f"{label}: m = {mm} tests against {100 * base:.1f}%")
        print(f"   p < 0.05 uncorrected {raw}   expected by chance {mm * 0.05:.1f}   "
              f"Bonferroni {int((tbl.p < 0.05 / mm).sum())}   "
              f"BH(q=0.05) {bh_survivors(tbl.p, 0.05)}   BH(q=0.10) {bh_survivors(tbl.p, 0.10)}")
        head = tbl.head(4).copy()
        head["pct"] = head.pct.map("{:.1f}".format)
        head["p"] = head.p.map("{:.4f}".format)
        print(head.to_string())
        print()
    print(f"league shot diet: {100 * rim_share_lg:.1f}% of shots from the restricted area")
    print("the four field-goal survivors, by diet:")
    for name in survivors_fg:
        print(f"  {name:<22} {100 * rim_share[name]:5.1f}% restricted area")

fg_raw, tp_raw = int((fg.p < 0.05).sum()), int((tp.p < 0.05).sum())
assert m_fg == 191 and m_tp == 124
assert len(shots) == 25000 and shots.PLAYER_NAME.nunique() == 547
assert abs(lg_fg - 0.47168) < 1e-9 and abs(lg_3p - 3607 / 9875) < 1e-9
assert fg_raw == 28 and tp_raw == 6
assert int((fg.p < 0.05 / m_fg).sum()) == 2
assert bh_survivors(fg.p, 0.05) == 4 and bh_survivors(fg.p, 0.10) == 6
assert int((tp.p < 0.05 / m_tp).sum()) == 0
assert bh_survivors(tp.p, 0.05) == 0 and bh_survivors(tp.p, 0.10) == 0
assert fg_raw > 2 * m_fg * 0.05                  # far more than chance would give
assert abs(tp_raw - m_tp * 0.05) <= 1            # indistinguishable from chance
assert fg.index[0] == "Nikola Jokic" and tp.index[0] == "Stephen Curry"
assert len(survivors_fg) == 4 and "Scoot Henderson" in survivors_fg
assert 0.29 < rim_share_lg < 0.30
assert all(rim_share[n] > rim_share_lg for n in survivors_fg)   # all four shoot nearer the rim
assert sum(rim_share[n] > 0.40 for n in survivors_fg) == 3
assert 0.33 < rim_share["Scoot Henderson"] < 0.34    # easy diet, worst percentage in the screen
assert float(fg.loc["Scoot Henderson", "pct"]) < 100 * lg_fg
assert 0.0048 < float(tp.p.iloc[0]) < 0.0050
assert float(fg.p.iloc[0]) < 0.0002

# --- 7. the exhibit: four p-value ladders against one BH line ----------------------
SERIES = [
    ("Team home record vs the league", screen.p.to_numpy(), sdt.sport_color("basketball")),
    ("All field goals, 50+ attempts", fg.p.to_numpy(), "#B23A3A"),
    ("Home-court edge, team by team", hca.p.to_numpy(), "#2C5E8A"),
    ("Three-point percentage, 30+ attempts", tp.p.to_numpy(), "#2E7D4F"),
]
fig, ax = plt.subplots(figsize=(8.8, 5.4))
xs = np.linspace(1e-3, 1, 200)
ax.plot(xs, ALPHA * xs, color="#20242B", lw=1.6, zorder=5,
        label=f"Benjamini-Hochberg line, q = {ALPHA}")
ax.axhline(ALPHA, color="#8A8577", lw=1.0, ls=":", zorder=2)
ax.text(0.985, ALPHA * 1.3, "uncorrected 0.05", fontsize=8, ha="right", color="#6C7079")
for label, pv, color in SERIES:
    pv = np.sort(pv)
    rank = np.arange(1, len(pv) + 1) / len(pv)
    keep = bh_survivors(pv, ALPHA)
    ax.plot(rank, np.clip(pv, 3e-7, 1.0), marker="o", ms=3.4, lw=0.9,
            color=color, alpha=0.85, zorder=4,
            label=f"{label} - m={len(pv)}, {keep} survive")
ax.set_yscale("log")
ax.set_ylim(3e-7, 3.0)
ax.set_xlim(0, 1.0)
ax.set_xlabel("rank of the p-value within its own screen (i / m)")
ax.set_ylabel("p-value (log scale)")
ax.set_title("Four screens, one threshold: below the line survives")
ax.legend(loc="lower right", fontsize=8.5, frameon=False)
ax.grid(True, which="major", axis="both")
sdt.save_fig(fig, "pvalue_ladders",
             source="Basketball-Reference 2023-24 schedule and the NBA_Shots_04_25 sample",
             asof="June 2026 bundle")

assert (np.sort(screen.p.to_numpy())[:bh_hits] <= ALPHA * np.arange(1, bh_hits + 1) / m).all()
assert not (np.sort(hca.p.to_numpy()) <= ALPHA * np.arange(1, m + 1) / m).any()
assert len(SERIES) == 4
print("\nall asserts passed")