"""
Tutorial 86 - Regression to the mean: why hot starts cool off.

Reads the bundled nba_home_results.csv (every 2023-24 final score), so it runs
offline. The whole tutorial is one deterministic experiment: split each team's
season into its first 41 and last 41 games, then regress second-half win
percentage on first-half win percentage. The fitted slope comes out at 0.77 -
below 1 - which is regression to the mean in a single number: a team one win
above average in the first half is, on average, only 0.77 wins above average
in the second. We show the fade at both extremes (the top five first-half
teams drop, the bottom five climb), rebuild the slope from first principles as
signal variance over total variance (the shrinkage estimate, 0.82 by a
completely different route), use it to out-predict the naive "second half =
first half" forecast, and finish with the curve that explains every October
overreaction: the fewer games you've seen, the harder the regression. No
randomness anywhere - every number recomputes identically on every run.

Run:  python downloads/86_regression_to_the_mean.py
"""
import os

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

import sdt_common as sdt

sdt.init("regression-to-the-mean")
HERE = os.path.dirname(os.path.abspath(__file__))

# --- one row per team per game, in schedule order ---------------------------------
games = pd.read_csv(os.path.join(HERE, "nba_home_results.csv"))
long = pd.concat([
    pd.DataFrame({"date": games["date"], "team": games["home_team"],
                  "win": (games["home_pts"] > games["away_pts"]).astype(int)}),
    pd.DataFrame({"date": games["date"], "team": games["away_team"],
                  "win": (games["away_pts"] > games["home_pts"]).astype(int)}),
])

HALF = 41


def half_split(first_n, last_n=None):
    """Per-team win%% over the first `first_n` games and the rest (or last `last_n`)."""
    rows = []
    for team, sub in long.groupby("team"):
        sub = sub.sort_values("date")
        early = sub["win"].iloc[:first_n].mean()
        late = (sub["win"].iloc[-last_n:] if last_n else sub["win"].iloc[first_n:]).mean()
        rows.append((team, early, late))
    return pd.DataFrame(rows, columns=["team", "first", "second"]).set_index("team")


halves = half_split(HALF, HALF)
x = halves["first"].to_numpy()
y = halves["second"].to_numpy()

with sdt.snippet("halves"):
    print(f"{len(games):,} games -> {len(long):,} team-games -> {len(halves)} teams\n")
    show = halves.copy()
    show["change"] = show["second"] - show["first"]
    print("best first halves:")
    print(show.sort_values("first", ascending=False).head(5).round(3).to_string())
    print("\nworst first halves:")
    print(show.sort_values("first").head(5).round(3).to_string())
    print(f"\nleague mean win%: first half {x.mean():.3f}, second half {y.mean():.3f}")

# --- the regression: second half on first half ------------------------------------
slope, intercept = np.polyfit(x, y, 1)
r = float(np.corrcoef(x, y)[0, 1])
assert 0.70 < slope < 0.85, f"slope moved: {slope}"
assert slope < 1.0, "the whole tutorial rests on this"

with sdt.snippet("slope"):
    print("regress second-half win% on first-half win% (30 teams):\n")
    print(f"  slope      {slope:.3f}")
    print(f"  intercept  {intercept:.3f}")
    print(f"  r          {r:.3f}")
    print()
    print("slope < 1 is regression to the mean, stated as arithmetic: a team")
    print(f"1 win above average before the break projects {slope:.2f} wins above")
    print("average after it. the other 0.23 wins were luck, and luck doesn't")
    print("carry over.")

# --- the fade at both extremes ----------------------------------------------------
top5 = halves.nlargest(5, "first")
bot5 = halves.nsmallest(5, "first")
assert top5["second"].mean() < top5["first"].mean()
assert bot5["second"].mean() > bot5["first"].mean()

with sdt.snippet("extremes"):
    print("the five best first-half teams:")
    print(f"  first half {top5['first'].mean():.3f}  ->  second half {top5['second'].mean():.3f}"
          f"   ({top5['first'].mean() * HALF:.1f} -> {top5['second'].mean() * HALF:.1f} wins per 41)")
    print("the five worst first-half teams:")
    print(f"  first half {bot5['first'].mean():.3f}  ->  second half {bot5['second'].mean():.3f}"
          f"   ({bot5['first'].mean() * HALF:.1f} -> {bot5['second'].mean() * HALF:.1f} wins per 41)")
    print()
    print("both extremes moved toward the middle - no coaching change required.")
    print("extreme records are where luck piled up in one direction, and the")
    print("luck resets to zero while the talent stays.")

# --- the shrinkage estimate: rebuild the slope from variance accounting -----------
var_obs = x.var(ddof=1)                     # spread of first-half win%
var_luck = (x * (1.0 - x)).mean() / HALF    # binomial noise in a 41-game win%
var_talent = var_obs - var_luck
shrink = var_talent / var_obs
assert 0.75 < shrink < 0.90, f"shrink moved: {shrink}"

mean1 = x.mean()
pred_naive = x                              # "second half = first half"
pred_shrunk = mean1 + shrink * (x - mean1)  # shrink toward the league mean


def rmse(pred):
    return float(np.sqrt(np.mean((y - pred) ** 2)))


with sdt.snippet("shrinkage"):
    print("where does 0.77 come from? split the variance of first-half win%:\n")
    print(f"  observed variance                 {var_obs:.5f}")
    print(f"  binomial luck in 41 games         {var_luck:.5f}   (mean of p(1-p)/41)")
    print(f"  what's left: talent variance      {var_talent:.5f}")
    print(f"  talent share = shrinkage factor   {shrink:.3f}")
    print()
    print(f"two routes, same answer: fitted slope {slope:.3f}, variance route {shrink:.3f}.")
    print("a 41-game record is ~4/5 talent, ~1/5 luck - so shrink it ~1/5 of")
    print("the way back to the mean before you forecast with it:\n")
    print(f"  RMSE, naive 'repeat the first half':      {rmse(pred_naive):.4f}")
    print(f"  RMSE, shrunk toward the league mean:      {rmse(pred_shrunk):.4f}")
    ext = np.argsort(-np.abs(x - mean1))[:10]
    rmse_ext_naive = float(np.sqrt(np.mean((y[ext] - pred_naive[ext]) ** 2)))
    rmse_ext_shrunk = float(np.sqrt(np.mean((y[ext] - pred_shrunk[ext]) ** 2)))
    print(f"  ...on the 10 most extreme teams only:     {rmse_ext_naive:.4f} -> {rmse_ext_shrunk:.4f}")
    print()
    print("the gain lives exactly where the takes live: at the extremes.")

assert rmse(pred_shrunk) < rmse(pred_naive)

# --- the earlier you look, the harder the regression ------------------------------
splits = [5, 10, 15, 20, 25, 30, 35, 41, 50, 60]
early_slopes = []
for n in splits:
    h = half_split(n)                       # first n games vs all the rest
    early_slopes.append(float(np.polyfit(h["first"], h["second"], 1)[0]))
slope10 = early_slopes[splits.index(10)]

ten = half_split(10)
hot = ten[ten["first"] >= 0.7]

with sdt.snippet("early-records"):
    print("slope of (rest of season) on (first n games), by n:\n")
    print("   games seen   slope")
    for n, s in zip(splits, early_slopes):
        print(f"     {n:>4}        {s:.2f}")
    print()
    print("after 10 games the slope is ~0.6: nearly half of what a 10-game")
    print("record says is noise. the 2023-24 receipts:\n")
    print("teams that started 8-2 (win% .800):")
    print(hot.round(3).to_string())
    print(f"\ncollectively .800 through ten games, {hot['second'].mean():.3f} the rest")
    print("of the way. still good - regression pulls toward the mean, not past")
    print("it - but every single one cooled off.")

assert len(hot) == 5 and hot["second"].mean() < 0.7

# --- the chart --------------------------------------------------------------------
color = sdt.sport_color("basketball")
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11.2, 4.9))

lims = np.array([0.05, 0.85])
ax1.plot(lims, lims, ls="--", lw=1.2, color="#6C7079")
ax1.annotate("slope 1: 'they'll repeat it'", xy=(0.175, 0.135), fontsize=9,
             color="#4A4F58", rotation=38, rotation_mode="anchor")
ax1.plot(lims, intercept + slope * lims, lw=2.2, color=color)
ax1.annotate(f"what actually happens:\nslope {slope:.2f}", xy=(0.28, intercept + slope * 0.28),
             xytext=(0.08, 0.55), fontsize=9, color="#20242B",
             arrowprops=dict(arrowstyle="-", color="#6C7079", lw=0.8))
ax1.scatter(x, y, s=42, color=color, edgecolor="#FBF7EE", linewidth=0.8, zorder=3)
for team, dx, dy in (("Philadelphia 76ers", -0.01, -0.055),
                     ("San Antonio Spurs", -0.03, 0.045),
                     ("Boston Celtics", -0.13, -0.012)):
    ax1.annotate(team.split()[-1], xy=(halves.loc[team, "first"] + dx,
                 halves.loc[team, "second"] + dy), fontsize=8.5, color="#4A4F58")
ax1.set_xlim(*lims)
ax1.set_ylim(*lims)
ax1.set_aspect("equal")
ax1.grid(axis="x")
ax1.set_xlabel("first-half win% (games 1-41)")
ax1.set_ylabel("second-half win% (games 42-82)")
ax1.set_title("Second halves regress toward the mean")

ax2.axhline(1.0, ls="--", lw=1.2, color="#6C7079")
ax2.annotate("slope 1 = records fully persist", xy=(59, 1.015), fontsize=9,
             color="#4A4F58", ha="right")
ax2.plot(splits, early_slopes, lw=2.2, color=color, marker="o", markersize=5)
for n, label, tdy in ((10, f"10 games in:\nslope {slope10:.2f}", -0.30),
                      (41, f"half a season:\nslope {slope:.2f}", -0.28)):
    s_val = early_slopes[splits.index(n)]
    ax2.scatter([n], [s_val], s=52, color="#20242B", zorder=3)
    ax2.annotate(label, xy=(n, s_val), xytext=(n + 3, s_val + tdy), fontsize=9,
                 color="#20242B",
                 arrowprops=dict(arrowstyle="-", color="#6C7079", lw=0.8))
ax2.set_xlim(0, 63)
ax2.set_ylim(0, 1.1)
ax2.set_xlabel("games seen before forecasting the rest")
ax2.set_ylabel("slope: rest-of-season on record so far")
ax2.set_title("The earlier you look, the harder the fade")

fig.suptitle("Regression to the mean in the 2023-24 NBA: the slope is not 1",
             fontweight="bold")
fig.tight_layout(rect=(0, 0.015, 1, 1))
sdt.save_fig(fig, "halves_slope",
             source="Bundled 2023-24 NBA results; per-team half-season splits")

print("done")