"""
Tutorial 94 - Survival curves: how long NFL head coaches really last.

"How long does a head coach last?" sounds like an average. It is not, because
the data stops before the story does: the coaches still in charge today have
not finished their tenures, and they include most of the longest ones. Drop
them and you measure only the failures; count them as finished and you cut
the survivors short. Survival analysis is the toolkit built for exactly this -
right-censored durations - and this tutorial writes its two workhorses by hand:
the Kaplan-Meier estimator (with Greenwood's standard error) and the log-rank
test, then uses them on every NFL head-coaching spell in the nflverse games
table, 1999 through the 2026 schedule.

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

    python downloads/94_survival_curves_kaplan_meier.py

Data: nflverse games table (github.com/nflverse/nfldata), June 2026 snapshot,
bundled as nfl_games_coaches.csv.
"""
import difflib
import math
import os

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

import sdt_common as sdt

sdt.init("survival-curves-kaplan-meier")
HERE = os.path.dirname(os.path.abspath(__file__))
CSV = os.path.join(HERE, "nfl_games_coaches.csv")

FRANCHISE = {"STL": "LA", "SD": "LAC", "OAK": "LV"}   # relocations, one franchise each
DATA_END = 2025        # last season with results; the 2026 schedule only says who is in charge


# --- the two estimators, written out in full ----------------------------------------
def kaplan_meier(T, E):
    """Kaplan-Meier survival table with Greenwood's standard error.

    T = seasons observed, E = 1 if the spell ended, 0 if still running (censored).
    At each time with an ending: n = spells still at risk, d = spells that ended,
    S(t) = S(t-1) * (1 - d/n), var = S^2 * sum d / (n (n - d)).
    """
    T, E = np.asarray(T), np.asarray(E)
    rows, S, g = [], 1.0, 0.0
    for t in np.unique(T[E == 1]):
        n = int((T >= t).sum())
        d = int(((T == t) & (E == 1)).sum())
        S *= 1 - d / n
        g += d / (n * (n - d)) if n > d else 0.0
        rows.append((int(t), n, d, d / n, S, S * math.sqrt(g)))
    return pd.DataFrame(rows, columns=["t", "at_risk", "ended", "hazard", "S", "se"])


def km_at(table, t):
    """S(t) read off a KM table: the last step at or before t."""
    past = table[table.t <= t]
    return 1.0 if past.empty else float(past.S.iloc[-1])


def km_median(table):
    """First time the survival curve reaches 0.5 or below."""
    hit = table[table.S <= 0.5]
    return int(hit.t.iloc[0]) if len(hit) else None


def log_rank(T, E, group):
    """Two-group log-rank test: observed vs expected endings in group 1, 1 df."""
    T, E, group = np.asarray(T), np.asarray(E), np.asarray(group, dtype=bool)
    O = Ex = V = 0.0
    for t in np.unique(T[E == 1]):
        at_risk = T >= t
        n, n1 = at_risk.sum(), (at_risk & group).sum()
        d = ((T == t) & (E == 1)).sum()
        O += ((T == t) & (E == 1) & group).sum()
        Ex += d * n1 / n
        if n > 1:
            V += d * (n1 / n) * (1 - n1 / n) * (n - d) / (n - 1)
    chi2 = (O - Ex) ** 2 / V
    return O, Ex, chi2, math.erfc(math.sqrt(chi2 / 2))   # chi-square(1) upper tail


# --- 1. games -> team-games -> coaching spells ---------------------------------------
games = pd.read_csv(CSV)
sides = []
for me, opp in (("home", "away"), ("away", "home")):
    s = games[["game_id", "season", "game_type", "gameday",
               f"{me}_team", f"{me}_coach", f"{me}_score", f"{opp}_score"]].copy()
    s.columns = ["game_id", "season", "game_type", "gameday", "team", "coach", "pts", "opp_pts"]
    sides.append(s)
tg = pd.concat(sides, ignore_index=True)
tg["team"] = tg.team.replace(FRANCHISE)
tg = tg.sort_values(["team", "gameday", "game_id"]).reset_index(drop=True)
tg["spell"] = (tg.coach != tg.groupby("team").coach.shift()).groupby(tg.team).cumsum()
tg["opener"] = tg.gameday == tg.groupby(["team", "season"]).gameday.transform("min")

raw = (tg.groupby(["team", "spell"])
         .agg(coach=("coach", "first"), first=("season", "min"),
              last=("season", "max"), week1=("opener", "first"))
         .reset_index())


def merge_returns(sp, gap=2):
    """A coach who is back with the same team within a season never left:
    fold the stand-ins in between into his spell."""
    out_all, merged = [], []
    for _, s in sp.groupby("team", sort=True):
        out = []
        for r in s.sort_values("first").to_dict("records"):
            j = next((k for k in range(len(out) - 1, -1, -1)
                      if out[k]["coach"] == r["coach"]), None)
            if j is not None and r["first"] - out[j]["last"] <= gap:
                merged.append((r["team"], r["coach"], out[j]["last"], r["first"],
                               [o["coach"] for o in out[j + 1:]]))
                out[j]["last"] = r["last"]
                del out[j + 1:]
            else:
                out.append(r)
        out_all += out
    return pd.DataFrame(out_all), merged


spells, merged = merge_returns(raw)
spells["running"] = spells["last"] > DATA_END                    # on the 2026 schedule
spells["seasons"] = spells["last"].clip(upper=DATA_END) - spells["first"] + 1
spells["left_trunc"] = (spells["first"] == 1999) & spells.week1   # hired before the file starts
cohort = spells[spells.week1 & ~spells.left_trunc & (spells.seasons > 0)].copy()
cohort["ended"] = (~cohort.running).astype(int)

n_mid = int((~spells.week1).sum())
with sdt.snippet("spells"):
    print(f"{len(games)} games -> {len(tg)} team-games, {tg.coach.nunique()} coach names")
    print(f"{len(raw)} unbroken runs of one coach per franchise")
    for team, coach, a, b, stand in merged:
        print(f"  merged: {coach} ({team}) - absent after {a}, back in {b}; "
              f"stand-ins folded in: {', '.join(stand)}")
    print(f"{len(spells)} coaching spells after merging returns\n")
    print(f"  already in charge at the file's first game (1999) : {int(spells.left_trunc.sum())}")
    print(f"  took over mid-season (interim starts)            : {n_mid}")
    print(f"  hired for 2026, no season observed yet           : "
          f"{int((spells.seasons <= 0).sum())}")
    print(f"  -> hires who opened a season, 2000-2025          : {len(cohort)}")
    print(f"       ended by the 2026 schedule                  : {int(cohort.ended.sum())}")
    print(f"       still in charge (right-censored)            : "
          f"{int((cohort.ended == 0).sum())}\n")
    ne = cohort[cohort.team == "NE"][["team", "coach", "first", "last", "seasons", "ended"]]
    print(ne.to_string(index=False))

assert len(games) == 7548 and len(tg) == 15096
assert tg.coach.notna().all() and tg.coach.nunique() == 177
assert tg.team.nunique() == 32                                   # relocations folded
names = sorted(tg.coach.unique())
assert not any(difflib.SequenceMatcher(None, a, b).ratio() > 0.85
               for i, a in enumerate(names) for b in names[i + 1:])  # no spelling variants
assert (tg[tg.season == 2026].groupby("team").coach.nunique() == 1).all()
assert len(raw) == 253 and len(spells) == 250
assert len(merged) == 1 and merged[0][:4] == ("NO", "Sean Payton", 2011, 2013)
assert merged[0][4] == ["Aaron Kromer", "Joe Vitt"]
assert int(spells.left_trunc.sum()) == 31 and n_mid == 42
assert int((spells.seasons <= 0).sum()) == 7
assert int(spells.left_trunc.sum()) + n_mid + 7 + len(cohort) == len(spells)   # nothing lost
assert len(cohort) == 170 and int(cohort.ended.sum()) == 145
assert int((cohort.ended == 0).sum()) == 25
payton = cohort[(cohort.team == "NO") & (cohort.coach == "Sean Payton")]
assert len(payton) == 1 and payton.seasons.iloc[0] == 16 and payton.ended.iloc[0] == 1
assert ne.coach.tolist() == ["Bill Belichick", "Jerod Mayo", "Mike Vrabel"]
assert ne.seasons.tolist() == [24, 1, 1] and ne.ended.tolist() == [1, 1, 0]
assert cohort.seasons.max() == 24 and cohort["first"].min() >= 2000
spans = sorted(c for o, n in FRANCHISE.items()                   # spells a relocation would split
               for c in set(games.loc[games.home_team == o, "home_coach"])
               & set(games.loc[games.home_team == n, "home_coach"]))
assert spans == ["Jeff Fisher", "Jon Gruden"]
assert int(((cohort["first"] == 2025) & (cohort.ended == 0)).sum()) == 6   # censored after one season

# --- 2. the two naive answers -------------------------------------------------------
done = cohort[cohort.ended == 1]
cohort["era"] = np.where(cohort["first"] <= 2012, "hired 2000-12", "hired 2013-25")
with sdt.snippet("naive"):
    print("completed spells only    : mean %.2f, median %.1f seasons  (n=%d)"
          % (done.seasons.mean(), done.seasons.median(), len(done)))
    print("running counted as ended : mean %.2f, median %.1f seasons  (n=%d)\n"
          % (cohort.seasons.mean(), cohort.seasons.median(), len(cohort)))
    for era, c in cohort.groupby("era"):
        cd = c[c.ended == 1]
        print(f"{era}: {len(c)} hires, {int((c.ended == 0).sum())} still running; "
              f"completed spells average {cd.seasons.mean():.2f} seasons "
              f"(median {cd.seasons.median():.0f})")
    old_mean = done[done["first"] <= 2012].seasons.mean()
    new_mean = done[done["first"] >= 2013].seasons.mean()
    print(f"naive verdict: recent spells {100 * (1 - new_mean / old_mean):.1f}% shorter")
    longest = cohort[cohort.ended == 0].sort_values(["seasons", "first"], ascending=[False, True])
    print("\nthe longest spells are the ones that have not ended:")
    print(longest.head(5)[["team", "coach", "first", "seasons"]].to_string(index=False))

assert len(done) == 145 and done.seasons.median() == 3.0
assert 4.165 < done.seasons.mean() < 4.175
assert 4.145 < cohort.seasons.mean() < 4.155 and cohort.seasons.median() == 3.0
assert int((cohort.era == "hired 2000-12").sum()) == 80
assert int(((cohort.era == "hired 2000-12") & (cohort.ended == 0)).sum()) == 0   # fully observed
assert int((cohort.era == "hired 2013-25").sum()) == 90
assert 5.095 < old_mean < 5.105 and 3.025 < new_mean < 3.035
assert 40.0 < 100 * (1 - new_mean / old_mean) < 41.0
assert longest.coach.iloc[0] == "Andy Reid" and longest.seasons.iloc[0] == 13
assert longest.seasons.head(4).tolist() == [13, 9, 9, 9]

# --- 3. Kaplan-Meier by hand ---------------------------------------------------------
km = kaplan_meier(cohort.seasons, cohort.ended)
km["lo"] = (km.S - 1.96 * km.se).clip(lower=0)
km["hi"] = (km.S + 1.96 * km.se).clip(upper=1)
med = km_median(km)
with sdt.snippet("km"):
    show = km.head(10).copy()
    for c in ("hazard", "S", "lo", "hi"):
        show[c] = show[c].map("{:.3f}".format)
    print(show[["t", "at_risk", "ended", "hazard", "S", "lo", "hi"]].to_string(index=False))
    print(f"...\nmedian tenure (first season S drops to 0.5 or below): {med} seasons")
    for t in (1, 3, 5, 10):
        print(f"  still in charge after {t:>2} season(s): {100 * km_at(km, t):.1f}%")

assert km[["at_risk", "ended"]].head(4).values.tolist() == [[170, 20], [144, 28], [111, 32], [75, 25]]
assert abs(km_at(km, 1) - 150 / 170) < 1e-12
assert abs(km.se.iloc[0] - math.sqrt((150 / 170) * (20 / 170) / 170)) < 1e-12  # Greenwood = binomial at step 1
assert 0.5055 < km_at(km, 3) < 0.5065 and med == 4
assert 0.2595 < km_at(km, 5) < 0.2605 and 0.0985 < km_at(km, 10) < 0.0995
assert list(km.hazard.head(4)) == sorted(km.hazard.head(4))        # risk climbs through year 4
assert int(km.ended.sum()) == 145
for i in range(len(km) - 1):                                        # risk-set bookkeeping
    t0, t1 = km.t.iloc[i], km.t.iloc[i + 1]
    cens = int(((cohort.seasons >= t0) & (cohort.seasons < t1) & (cohort.ended == 0)).sum())
    assert km.at_risk.iloc[i + 1] == km.at_risk.iloc[i] - km.ended.iloc[i] - cens
assert ((km.lo <= km.S) & (km.S <= km.hi)).all()
assert abs(km.hazard.iloc[0] - 20 / 170) < 1e-12 and abs(km.hazard.iloc[3] - 25 / 75) < 1e-12
since23 = cohort[cohort["first"] >= 2023]
assert not (kaplan_meier(since23.seasons, since23.ended).S <= 0.5).any()  # median not reached

# --- 4. has the leash shortened? two eras, one log-rank test -------------------------
recent = (cohort.era == "hired 2013-25").to_numpy()
O, Ex, chi2, p_era = log_rank(cohort.seasons, cohort.ended, recent)
km_old = kaplan_meier(cohort.seasons[~recent], cohort.ended[~recent])
km_new = kaplan_meier(cohort.seasons[recent], cohort.ended[recent])

rng = np.random.default_rng(94)
PERMS = 2000
T_arr, E_arr = cohort.seasons.to_numpy(), cohort.ended.to_numpy()
perm_chi = np.array([log_rank(T_arr, E_arr, rng.permutation(recent))[2] for _ in range(PERMS)])
p_perm = float((perm_chi >= chi2).mean())

with sdt.snippet("logrank"):
    for label, tab, mask in (("hired 2000-12", km_old, ~recent), ("hired 2013-25", km_new, recent)):
        print(f"{label}: n={int(mask.sum())}  median {km_median(tab)}  "
              f"S(2)={km_at(tab, 2):.3f}  S(3)={km_at(tab, 3):.3f}  S(5)={km_at(tab, 5):.3f}")
    print(f"\nlog-rank, recent hires: observed endings {O:.0f}, expected {Ex:.2f}")
    print(f"chi-square = {chi2:.3f} on 1 df, p = {p_era:.3f}")
    print(f"permutation check ({PERMS} label shuffles): p = {p_perm:.3f}")

old_T = cohort.seasons[~recent]
for t in km_old.t:                          # no censoring: KM is the plain empirical survival
    assert abs(km_at(km_old, t) - float((old_T > t).mean())) < 1e-12
assert km_median(km_old) == 4 and km_median(km_new) == 3
assert abs(km_at(km_old, 3) - 43 / 80) < 1e-12 and 0.4745 < km_at(km_new, 3) < 0.4755   # uncensored: exact
assert abs(km_at(km_old, 2) - 59 / 80) < 1e-12 and 0.6855 < km_at(km_new, 2) < 0.6865
assert int(O) == 65 and 61.26 < Ex < 61.28
assert 0.5415 < chi2 < 0.5425 and 0.4605 < p_era < 0.4615
assert abs(log_rank(T_arr, E_arr, ~recent)[2] - chi2) < 1e-9          # symmetric in the groups
assert abs(math.erfc(math.sqrt(3.841459 / 2)) - 0.05) < 1e-6           # chi-square(1) critical value
assert abs(p_perm - p_era) < 0.03 and 0.44 < p_perm < 0.48

# --- 5. a split that does separate: the first-season record, landmarked ---------------
reg = tg[(tg.game_type == "REG") & tg.pts.notna()].copy()
reg["w"] = (reg.pts > reg.opp_pts) + 0.5 * (reg.pts == reg.opp_pts)
first_year = reg.merge(cohort[["team", "coach", "first"]],
                       left_on=["team", "coach", "season"], right_on=["team", "coach", "first"])
rec = first_year.groupby(["team", "coach", "first"]).w.agg(["sum", "size"]).reset_index()
cohort = cohort.merge(rec, on=["team", "coach", "first"], how="left")
cohort["win1"] = (cohort["sum"] / cohort["size"]) > 0.5

land = cohort[cohort.seasons >= 2]                          # reached season 2 on record
O2, Ex2, chi2_land, p_land = log_rank(land.seasons, land.ended, land.win1)
km_win = kaplan_meier(land.seasons[land.win1], land.ended[land.win1])
km_lose = kaplan_meier(land.seasons[~land.win1], land.ended[~land.win1])
_, _, chi2_naive, p_naive = log_rank(cohort.seasons, cohort.ended, cohort.win1)
one_and_done = cohort[(cohort.seasons == 1) & (cohort.ended == 1)]

with sdt.snippet("landmark"):
    print(f"hires observed into a second season: {len(land)} "
          f"({int(land.win1.sum())} had a winning first season)\n")
    for label, tab, mask in (("winning first season", km_win, land.win1),
                             (".500 or worse", km_lose, ~land.win1)):
        print(f"{label:<21} n={int(mask.sum()):>3}  ended {int(land.ended[mask].sum()):>3}  "
              f"median {km_median(tab)}  S(3)={km_at(tab, 3):.3f}  S(5)={km_at(tab, 5):.3f}")
    print(f"\nlog-rank from the season-2 landmark: chi-square = {chi2_land:.3f}, p = {p_land:.3f}")
    print(f"the wrong version, all {len(cohort)} hires split from day one: "
          f"chi-square = {chi2_naive:.2f}, p = {p_naive:.2g}")
    print(f"  (it counts {len(one_and_done)} one-season spells, "
          f"{int((~one_and_done.win1).sum())} of them with a .500-or-worse record, as evidence)")

assert cohort["size"].notna().all() and cohort["size"].between(6, 17).all()
assert len(land) == 144 and int(land.win1.sum()) == 57
assert int(land.ended[land.win1].sum()) == 46 and int(land.ended[~land.win1].sum()) == 79
assert km_median(km_win) == 4 and km_median(km_lose) == 3
assert 0.7075 < km_at(km_win, 3) < 0.7085 and 0.4855 < km_at(km_lose, 3) < 0.4865
assert 0.4095 < km_at(km_win, 5) < 0.4105 and 0.2205 < km_at(km_lose, 5) < 0.2215
assert 4.5275 < chi2_land < 4.5285 and 0.0325 < p_land < 0.0335
assert O2 < Ex2                                   # winners end less often than expected
assert 0.025 < p_land < 0.05                      # clears 0.05, not a two-test Bonferroni 0.025
assert 9.905 < chi2_naive < 9.915 and p_naive < 0.002 and p_naive < p_land
assert len(one_and_done) == 20 and int((~one_and_done.win1).sum()) == 19
losers = land[~land.win1]
assert ((losers["sum"] == 8) & (losers["size"] == 17)).any()          # an 8-9 first season
assert ((losers["sum"] == 1) & (losers["size"] == 16)).any()          # a 1-15 first season
assert all(km_at(km_win, t) > km_at(km_lose, t) for t in range(2, 7))
assert km_at(km_win, 7) < km_at(km_lose, 7)                           # the curves cross at seven
assert chi2_naive > 2 * chi2_land

# --- 6. the exhibit ------------------------------------------------------------------
BROWN, BLUE, GREY = sdt.sport_color("football"), "#2C5E8A", "#8A8577"


def step(ax, tab, **kw):
    """Draw a KM table as a right-continuous step curve starting at S=1."""
    xs, ys = [0], [1.0]
    prev = 1.0
    for t, s_val in zip(tab.t, tab.S):
        xs += [t, t]
        ys += [prev, s_val]
        prev = s_val
    ax.plot(xs, ys, **kw)


fig, (a1, a2) = plt.subplots(1, 2, figsize=(10.4, 4.8), sharey=True)
step(a1, km_old, color=BROWN, lw=2.0, label=f"Hired 2000-12 (n={int((~recent).sum())})")
step(a1, km_new, color=BLUE, lw=2.0, label=f"Hired 2013-25 (n={int(recent.sum())}), Kaplan-Meier")
ended_new = recent & (cohort.ended == 1).to_numpy()
naive_new = kaplan_meier(cohort.seasons[ended_new], cohort.ended[ended_new])
step(a1, naive_new, color=BLUE, lw=1.4, ls="--", label="Hired 2013-25, running spells dropped")
assert all(km_at(naive_new, t) < km_at(km_new, t) for t in range(1, 9))   # dropping censored rows biases low
a1.axhline(0.5, color=GREY, lw=0.8, ls=":")
a1.set_xlim(0, 16)
a1.set_ylim(0, 1.02)
a1.set_xlabel("seasons since hire")
a1.set_ylabel("share of spells still running")
a1.set_title(f"Two eras: log-rank p = {p_era:.2f}", fontsize=11.5)
a1.legend(loc="upper right", fontsize=8.2, frameon=False)
a1.grid(True, axis="both")

step(a2, km_win, color="#2E7D4F", lw=2.0, label=f"Winning first season (n={int(land.win1.sum())})")
step(a2, km_lose, color="#B23A3A", lw=2.0, label=f".500 or worse (n={int((~land.win1).sum())})")
a2.axhline(0.5, color=GREY, lw=0.8, ls=":")
a2.set_xlim(0, 16)
a2.set_xlabel("seasons since hire (from the season-2 landmark)")
a2.set_title(f"First-season record: log-rank p = {p_land:.2f}", fontsize=11.5)
a2.legend(loc="upper right", fontsize=8.2, frameon=False)
a2.grid(True, axis="both")
fig.tight_layout()
sdt.save_fig(fig, "tenure_curves", source="nflverse games table (github.com/nflverse/nfldata)",
             asof="June 2026")
print("\nall asserts passed")