"""
Tutorial 100 - Gaussian mixture models and the EM algorithm, from scratch: the court's zones from coordinates alone.

24,947 half-court shots from the 2023-24 NBA season carry two numbers each, where the shot
was taken. A Gaussian mixture says those points were drawn from a few overlapping bell
shapes, and the EM algorithm finds the shapes by alternating two easy steps: share every
shot out among the current components (E), then refit each component to its share (M).

Six components, never shown the league's zone labels, agree with them on 85.4% of shots,
where k-means with the same six manages 67.3%: the corner threes come out as two strips
half a foot wide that round k-means clusters cannot draw. The lesson is in the likelihood
itself. 704 shots are logged at exactly the middle of the basket, and an eight-component
fit pulls one component onto that single point, where the likelihood has no maximum. With
a variance floor the fit is finite, but the Bayesian information criterion still spends a
component on that one spot and is still asking for more components at ten.

Written by hand: the one-dimensional EM on paper, the two-dimensional EM with full
covariances and the log-sum-exp trick, random restarts, the collapse and its floor, BIC,
responsibilities, and a k-means comparison. No scikit-learn, no scipy. Every number the
tutorial page quotes is asserted below.

Run:

    python downloads/100_gaussian_mixture_models_em_from_scratch.py

Data: the public NBA_Shots_04_25 shot log (github.com/DomSamangy/NBA_Shots_04_25),
2023-24 season, a 25,000-row random sample (random_state=42) bundled as nba_league_shots.csv.
The court dimension quoted is from the NBA rulebook, Rule No. 1, Section I.
"""
import os
import sys

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.patches import Arc, Ellipse

import sdt_common as sdt

try:
    sys.stdout.reconfigure(encoding="utf-8")
except (AttributeError, ValueError):
    pass

sdt.init("gaussian-mixture-models-em-from-scratch")
os.chdir(os.path.dirname(os.path.abspath(__file__)))       # the CSV sits beside this script


def near(a, b, tol=5e-4):
    return abs(a - b) < tol


# ================================================================ 1. the court in feet ==
with sdt.snippet("data"):
    import numpy as np
    import pandas as pd

    shots = pd.read_csv("nba_league_shots.csv")
    r = np.hypot(shots.LOC_X, shots.LOC_Y - 5.25)            # feet from the middle of the basket
    print(f"{len(shots):,} shots; floor(r) equals SHOT_DISTANCE on {(np.floor(r) == shots.SHOT_DISTANCE).mean():.3%}")

    shots = shots[shots.BASIC_ZONE != "Backcourt"].reset_index(drop=True)
    X = shots[["LOC_X", "LOC_Y"]].to_numpy(float)
    dist = np.hypot(X[:, 0], X[:, 1] - 5.25)
    print(f"{len(shots):,} half-court shots, |x| at most {np.abs(X[:, 0]).max()} ft")
    for zone, n in shots.BASIC_ZONE.value_counts().items():
        print(f"  {zone:22s} {n:5,d}")

    three = (shots.SHOT_TYPE == "3PT Field Goal").to_numpy()
    low = X[:, 1] < 14                                       # the straight part of the line
    print(f"threes below y = 14: none closer than {np.abs(X[three & low, 0]).min()} ft to the middle line;"
          f" threes above: none closer than {dist[three & ~low].min():.2f} ft to the basket")
    print("threes the zone column calls Mid-Range:", int((three & (shots.BASIC_ZONE == "Mid-Range")).sum()))

    spots = shots.groupby(["LOC_X", "LOC_Y"]).size().sort_values(ascending=False)
    at_rim = (shots.LOC_X == 0) & (shots.LOC_Y == 5.25)
    print(f"busiest spot: x = 0, y = 5.25, {spots.iloc[0]} shots ({at_rim.mean():.2%});"
          f" next busiest: {spots.iloc[1]}")

n_shots = len(shots)
assert n_shots == 24947 and int((np.floor(r) == pd.read_csv("nba_league_shots.csv").SHOT_DISTANCE).sum()) == 24973
assert f"{(np.floor(r) == pd.read_csv('nba_league_shots.csv').SHOT_DISTANCE).mean():.3%}" == "99.892%"
assert np.abs(X[:, 0]).max() == 25.0
zc = shots.BASIC_ZONE.value_counts()
assert (zc["Restricted Area"], zc["Above the Break 3"], zc["In The Paint (Non-RA)"], zc["Mid-Range"],
        zc["Left Corner 3"], zc["Right Corner 3"]) == (7379, 7255, 4971, 2779, 1310, 1253)
corner = shots.BASIC_ZONE.str.contains("Corner").to_numpy()
assert np.abs(X[three & low, 0]).min() == 21.9 and f"{dist[three & ~low].min():.2f}" == "23.77"
assert int((three & (shots.BASIC_ZONE == "Mid-Range")).sum()) == 4 and X[corner, 1].max() < 14
assert spots.iloc[0] == 704 and int(at_rim.sum()) == 704 and spots.index[0] == (0.0, 5.25)
assert spots.iloc[1] == 22 and f"{at_rim.mean():.2%}" == "2.82%"
assert (shots.SHOT_DISTANCE[at_rim] == 0).all() and (shots.BASIC_ZONE[at_rim] == "Restricted Area").all()
assert set(np.round(X[:, 0] * 10) / 10 == X[:, 0]) == {True}       # x is recorded to a tenth of a foot
raw = pd.read_csv("nba_league_shots.csv")
r_raw = np.hypot(raw.LOC_X, raw.LOC_Y - 5.25)
miss = np.floor(r_raw) != raw.SHOT_DISTANCE
assert int(miss.sum()) == 27 and ((r_raw - raw.SHOT_DISTANCE)[miss].abs() < 1e-9).all()   # all a float hair below a whole foot
assert len(raw) - n_shots == 53
assert corner.sum() == 2563 and (three[corner]).all()


# ================================================================ 2. EM on paper, one dimension ==
with sdt.snippet("oned"):
    def normal_pdf(x, mean, sd):
        return np.exp(-0.5 * ((x - mean) / sd) ** 2) / (sd * np.sqrt(2 * np.pi))

    def em_1d(x, w, mean, sd, tol=1e-8, max_iter=1000):
        history = []
        for it in range(max_iter):
            dens = w * normal_pdf(x[:, None], mean, sd)      # E-step: weight x density...
            total = dens.sum(axis=1)
            history.append(np.log(total).sum())            # (the log-likelihood so far)
            resp = dens / total[:, None]                    # ...as each component's share
            nk = resp.sum(axis=0)                           # M-step: shots per component,
            w = nk / len(x)                                 # its weight,
            mean = (resp * x[:, None]).sum(axis=0) / nk     # its mean
            sd = np.sqrt((resp * (x[:, None] - mean) ** 2).sum(axis=0) / nk)   # and its spread
            if it and history[-1] - history[-2] < tol:
                break
        return w, mean, sd, history

    w1, m1, s1, hist = em_1d(dist, np.array([0.5, 0.5]), np.array([5.0, 25.0]), np.array([5.0, 5.0]))
    print(f"{len(hist)} rounds; log-likelihood {hist[0]:,.1f} -> {hist[1]:,.1f} -> {hist[2]:,.1f}"
          f" ... {hist[-1]:,.2f}; it never fell: {bool(np.all(np.diff(hist) >= 0))}")
    print(f"near: weight {w1[0]:.4f}, mean {m1[0]:.3f} ft, sd {s1[0]:.3f}")
    print(f"far:  weight {w1[1]:.4f}, mean {m1[1]:.3f} ft, sd {s1[1]:.3f}")

    parts = w1 * normal_pdf(22.0, m1, s1)                   # one shot from 22 feet
    print(f"22 ft: near {parts[0]:.6f}, far {parts[1]:.6f}, far share {parts[1] / parts.sum():.4f}")

    grid = np.arange(0, 80, 0.001)
    lead = np.sign(w1[1] * normal_pdf(grid, m1[1], s1[1]) - w1[0] * normal_pdf(grid, m1[0], s1[0]))
    flips = grid[1:][np.diff(lead) != 0]
    inside = (dist > flips[0]) & (dist < flips[1])
    print(f"the far component wins from {flips[0]:.2f} to {flips[1]:.2f} ft:"
          f" {inside.sum():,} shots, {three[inside].mean():.1%} of them threes")
    print(f"shots beyond {flips[1]:.2f} ft, handed back to the near component: {(dist > flips[1]).sum()}")

assert len(hist) == 39 and np.all(np.diff(hist) >= 0)
assert near(hist[0], -86644.957, 1e-2) and near(hist[-1], -83799.334, 1e-2)
assert f"{w1[0]:.4f}" == "0.5924" and f"{m1[0]:.3f}" == "6.065" and f"{s1[0]:.3f}" == "5.444"
assert f"{w1[1]:.4f}" == "0.4076" and f"{m1[1]:.3f}" == "25.361" and f"{s1[1]:.3f}" == "1.984"
assert near(parts[1] / parts.sum(), 0.9703, 5e-5)
# the hand version of the 22-foot share, from the four-decimal parameters the page prints
hand_near = 0.5924 * np.exp(-0.5 * ((22 - 6.065) / 5.444) ** 2) / (5.444 * np.sqrt(2 * np.pi))
hand_far = 0.4076 * np.exp(-0.5 * ((22 - 25.361) / 1.984) ** 2) / (1.984 * np.sqrt(2 * np.pi))
assert near(hand_far / (hand_near + hand_far), parts[1] / parts.sum(), 5e-4)
# the page's worked example, step by step, at the precision it is written
z_near, z_far = (22 - 6.065) / 5.444, (22 - 25.361) / 1.984
assert f"{z_near:.4f}" == "2.9271" and f"{z_near ** 2 / 2:.4f}" == "4.2839" and f"{np.exp(-z_near ** 2 / 2):.5f}" == "0.01379"
assert f"{5.444 * np.sqrt(2 * np.pi):.4f}" == "13.6461" and f"{0.01379 / 13.6461:.7f}" == "0.0010105"
assert f"{0.5924 * 0.0010105:.6f}" == "0.000599"
assert f"{-z_far:.4f}" == "1.6941" and f"{z_far ** 2 / 2:.4f}" == "1.4349" and f"{np.exp(-z_far ** 2 / 2):.5f}" == "0.23814"
assert f"{1.984 * np.sqrt(2 * np.pi):.4f}" == "4.9732" and f"{0.23814 / 4.9732:.6f}" == "0.047885"
assert f"{0.4076 * 0.047885:.6f}" == "0.019518" and f"{0.019518 / (0.019518 + 0.000599):.4f}" == "0.9702"
assert f"{flips[0]:.2f}" == "19.86" and f"{flips[1]:.2f}" == "36.78" and len(flips) == 2
assert int(inside.sum()) == 10179 and f"{three[inside].mean():.1%}" == "96.3%"
assert int((dist > flips[1]).sum()) == 19


# ================================================================ 3. two dimensions ==
with sdt.snippet("twod"):
    def log_normal_2d(X, mean, cov):
        """log density of every shot under every component: an (n, k) table."""
        a, b, c = cov[:, 0, 0], cov[:, 0, 1], cov[:, 1, 1]
        det = a * c - b * b
        dx = X[:, None, 0] - mean[None, :, 0]
        dy = X[:, None, 1] - mean[None, :, 1]
        q = (c * dx * dx - 2 * b * dx * dy + a * dy * dy) / det   # the 2 x 2 inverse, written out
        return -0.5 * q - np.log(2 * np.pi) - 0.5 * np.log(det)

    def gmm(X, k, seed, floor=0.0, tol=1e-3, max_iter=1000):
        rng = np.random.default_rng(seed)
        mean = X[rng.choice(len(X), k, replace=False)]        # start: k random shots as centres
        cov = np.repeat(np.cov(X.T, bias=True)[None], k, axis=0)
        w = np.full(k, 1 / k)
        history = []
        for it in range(1, max_iter + 1):
            L = np.log(w) + log_normal_2d(X, mean, cov)
            top = L.max(axis=1)
            lse = top + np.log(np.exp(L - top[:, None]).sum(axis=1))   # log-sum-exp, no underflow
            history.append(lse.sum())
            resp = np.exp(L - lse[:, None])                    # E-step
            nk = resp.sum(axis=0)                              # M-step
            d = X[:, None, :] - (resp.T @ X / nk[:, None])[None]
            new_cov = np.einsum("nk,nki,nkj->kij", resp, d, d) / nk[:, None, None] + floor * np.eye(2)
            det = new_cov[:, 0, 0] * new_cov[:, 1, 1] - new_cov[:, 0, 1] ** 2
            if det.min() < 1e-9:                               # a component has shrunk onto a point
                return dict(w=w, mean=mean, cov=cov, resp=resp, ll=history[-1], it=it,
                            history=history, collapsed=True)
            w, mean, cov = nk / len(X), resp.T @ X / nk[:, None], new_cov
            if it > 1 and history[-1] - history[-2] < tol:
                break
        return dict(w=w, mean=mean, cov=cov, resp=resp, ll=history[-1], it=it,
                    history=history, collapsed=False)

    fits = [gmm(X, 6, seed) for seed in range(10)]
    lls = np.array([f["ll"] for f in fits])
    print("ten starts, final log-likelihoods:", ", ".join(f"{v:,.1f}" for v in sorted(lls)))
    best = fits[int(lls.argmax())]
    order = np.lexsort((best["mean"][:, 0], best["mean"][:, 1]))   # one fixed order for the labels
    w6, mean6, cov6 = best["w"][order], best["mean"][order], best["cov"][order]
    label = best["resp"][:, order].argmax(axis=1)

    table = pd.crosstab(label, shots.BASIC_ZONE)
    print(" comp  weight   centre (x, y)   sd x   sd y   corr  largest zone            share  makes")
    for j in range(6):
        sx, sy = np.sqrt(cov6[j, 0, 0]), np.sqrt(cov6[j, 1, 1])
        top = table.loc[j].idxmax()
        print(f"  {j + 1}   {w6[j]:.3f}  ({mean6[j, 0]:6.2f},{mean6[j, 1]:6.2f})  {sx:5.2f}  {sy:5.2f}"
              f"  {cov6[j, 0, 1] / (sx * sy):+.2f}  {top:22s}  {table.loc[j, top] / table.loc[j].sum():.3f}"
              f"  {shots.SHOT_MADE[label == j].mean():.3f}")
    purity = table.max(axis=1).sum() / len(shots)
    print(f"shots in their component's largest zone: {purity:.1%}")
    worst = fits[int(lls.argmin())]
    print("corner strips (sd x under 0.6 ft) in the best fit:", int((np.sqrt(cov6[:, 0, 0]) < 0.6).sum()),
          "; in the lowest:", int((np.sqrt(worst["cov"][:, 0, 0]) < 0.6).sum()))

assert sorted(np.round(lls, 1))[-1] == round(best["ll"], 1)
assert near(best["ll"], -159642.34, 0.05)
n_best = int((lls > lls.max() - 0.5).sum())
assert n_best == 8 and near(sorted(lls)[0], -164167.4, 0.05) and near(sorted(lls)[1], -164167.4, 0.05)
assert f"{purity:.1%}" == "85.4%"
assert table.loc[0].idxmax() == "Restricted Area" and near(w6[0], 0.303, 5e-4)
corner_comps = [j for j in range(6) if np.sqrt(cov6[j, 0, 0]) < 0.6]
assert len(corner_comps) == 2
for j in corner_comps:
    zone = table.loc[j].idxmax()
    assert "Corner" in zone and table.loc[j, zone] / table.loc[j].sum() > 0.85
    assert 22.9 < abs(mean6[j, 0]) < 23.2 and np.sqrt(cov6[j, 0, 0]) > 0.45
left = [j for j in corner_comps if mean6[j, 0] > 0][0]
right = [j for j in corner_comps if mean6[j, 0] < 0][0]
assert table.loc[left].idxmax() == "Left Corner 3" and table.loc[right].idxmax() == "Right Corner 3"
abv = [j for j in range(6) if table.loc[j].idxmax() == "Above the Break 3"]
assert len(abv) == 2 and sorted(np.sign(mean6[abv, 0])) == [-1, 1]
mid = [j for j in range(6) if table.loc[j].idxmax() == "In The Paint (Non-RA)"]
assert len(mid) == 1
assert "Mid-Range" not in {table.loc[j].idxmax() for j in range(6)}
assert near(shots.SHOT_MADE[label == 0].mean(), 0.643, 5e-4)
mr_in_mid = table.loc[mid[0], "Mid-Range"] / zc["Mid-Range"]
assert f"{mr_in_mid:.1%}" == "94.4%" and mid == [3]
shares = [table.loc[j].max() / table.loc[j].sum() for j in range(6)]
assert [f"{v:.3f}" for v in shares] == ["0.930", "0.930", "0.870", "0.625", "0.984", "0.975"]
assert (left, right) == (1, 2) and table.loc[left, "Left Corner 3"] == 1310 and table.loc[right, "Right Corner 3"] == 1252
assert [f"{np.sqrt(cov6[j, 0, 0]):.2f}" for j in (1, 2)] == ["0.51", "0.55"]
assert [f"{np.sqrt(cov6[j, 1, 1]):.2f}" for j in (1, 2)] == ["3.33", "3.92"]
assert [f"{abs(mean6[j, 0]):.2f}" for j in (1, 2)] == ["23.11", "23.04"]
assert [f"{np.sqrt(cov6[0, i, i]):.2f}" for i in (0, 1)] == ["1.70", "1.16"]
assert [f"{np.sqrt(cov6[3, i, i]):.2f}" for i in (0, 1)] == ["7.50", "5.05"]
rho = [cov6[j, 0, 1] / np.sqrt(cov6[j, 0, 0] * cov6[j, 1, 1]) for j in (4, 5)]
assert [f"{v:+.2f}" for v in rho] == ["-0.86", "+0.81"]
wsd = np.sqrt(worst["cov"][:, [0, 1], [0, 1]])
flat = int(np.argmax(wsd[:, 0]))
assert f"{wsd[flat, 0]:.1f}" == "17.6" and f"{wsd[flat, 1]:.1f}" == "1.9" and wsd[flat, 1] < 2 < 6 < mean6[3, 1]
assert np.sqrt(worst["cov"][:, 0, 0]).min() > 0.6                    # no corner strip in the lower fit


# ================================================================ 4. the collapse ==
with sdt.snippet("collapse"):
    bad = gmm(X, 8, seed=0)
    h = np.array(bad["history"])
    print(f"8 components, start 0: collapsed = {bad['collapsed']} at round {bad['it']}")
    print("last rounds' log-likelihood:", ", ".join(f"{v:,.0f}" for v in h[-6:]))
    sds = np.sqrt(bad["cov"][:, [0, 1], [0, 1]])
    j = int(np.argmin(sds.prod(axis=1)))
    print(f"component {j}: weight x shots = {bad['w'][j] * len(X):.1f}, centre ({bad['mean'][j, 0]:.4f},"
          f" {bad['mean'][j, 1]:.4f}), sd ({sds[j, 0]:.2e}, {sds[j, 1]:.2e}) ft")
    owned = bad["resp"][:, j] > 0.5
    print(f"shots it owns: {owned.sum()}, {int(at_rim[owned].sum())} of them at the basket spot;"
          f" the nearest owned shot that is not: {np.hypot(*(X[owned & ~at_rim] - [0, 5.25]).T).min():.2f} ft away")
    print(f"every tenfold shrink of both sds adds 704 x 2 ln 10 = {704 * 2 * np.log(10):,.1f}")

    ok = gmm(X, 8, seed=0, floor=0.01)                         # a floor: no sd below 0.1 ft
    sds = np.sqrt(ok["cov"][:, [0, 1], [0, 1]])
    j = int(np.argmin(sds.prod(axis=1)))
    print(f"with the floor: collapsed = {ok['collapsed']}, log-likelihood {ok['ll']:,.1f};"
          f" tightest component weight x shots = {ok['w'][j] * len(X):.1f},"
          f" sd ({sds[j, 0]:.3f}, {sds[j, 1]:.3f})")

assert bad["collapsed"] and bad["it"] == 189
assert np.all(np.diff(h) > 0) and h[-1] - h[-6] > 2500
jb = int(np.argmin(np.sqrt(bad["cov"][:, [0, 1], [0, 1]]).prod(axis=1)))
assert near(bad["mean"][jb, 0], 0.0, 5e-3) and near(bad["mean"][jb, 1], 5.25, 5e-3)
assert near(bad["w"][jb] * len(X), 738.8, 0.05)
assert int((bad["resp"][:, jb] > 0.5).sum()) == 725 and int(at_rim[bad["resp"][:, jb] > 0.5].sum()) == 704
assert f"{704 * 2 * np.log(10):,.1f}" == "3,242.0"
assert f"{h[-5]:,.0f}" == "-157,197" and f"{h[-6]:,.0f}" == "-157,462" and f"{h[-1]:,.0f}" == "-154,691"
assert f"{bad['mean'][jb, 0]:.4f}" == "0.0004" and f"{bad['mean'][jb, 1]:.4f}" == "5.2520"
sdb = np.sqrt(bad["cov"][jb, [0, 1], [0, 1]])
assert f"{sdb[0]:.3f}" == "0.012" and f"{sdb[1]:.3f}" == "0.028"
assert not ok["collapsed"]
jo = int(np.argmin(np.sqrt(ok["cov"][:, [0, 1], [0, 1]]).prod(axis=1)))
sd_ok = np.sqrt(ok["cov"][jo, [0, 1], [0, 1]])
assert np.all(sd_ok > 0.1) and np.all(sd_ok < 0.13)
assert near(ok["mean"][jo, 0], 0.0, 0.01) and near(ok["mean"][jo, 1], 5.26, 0.01)
assert near(ok["ll"], -156918.6, 0.05) and f"{ok['w'][jo] * len(X):.1f}" == "824.2"
assert [f"{v:.3f}" for v in sd_ok] == ["0.115", "0.120"]


# ================================================================ 5. how many components ==
with sdt.snippet("bic"):
    def bic(fit, k, n):
        return -2 * fit["ll"] + (6 * k - 1) * np.log(n)        # 6k - 1 free numbers

    sweep = []
    for k in range(1, 11):
        runs = [gmm(X, k, seed, floor=0.01) for seed in range(4)]
        top = max(runs, key=lambda f: f["ll"])
        sds = np.sqrt(top["cov"][:, [0, 1], [0, 1]]).prod(axis=1)
        spot = bool((sds < 0.02).any())                        # a component sitting on one spot?
        sweep.append((k, top["ll"], bic(top, k, len(X)), spot))
        print(f"k = {k:2d}  log-likelihood {top['ll']:12,.1f}  BIC {sweep[-1][2]:11,.1f}"
              f"  on the basket spot: {'yes' if spot else 'no'}")
    drops = np.diff([s[2] for s in sweep])
    print(f"BIC falls at every step; smallest fall {-drops.max():,.1f} (k = 9 to 10)")

assert len(sweep) == 10
assert all(drops < 0)
assert int(np.argmax(drops)) == 8
assert [s[3] for s in sweep] == [False] * 7 + [True] * 3
assert near(sweep[5][1], -159643.8, 0.1) and near(sweep[9][2], 310998.3, 0.1)
assert near(sweep[0][2], 374164.4, 0.1) and near(sweep[7][2], 314313.1, 0.1) and near(sweep[6][2], 317911.8, 0.1)
assert f"{-drops.max():,.1f}" == "430.6" and near(sweep[7][1], ok["ll"], 1e-6)   # k = 8 best start is start 0
assert sweep[1][2] - sweep[0][2] < -19000 and near(sweep[5][2], 319641.9, 0.1)
bic_gain_8 = sweep[6][2] - sweep[7][2]
assert f"{bic_gain_8:,.1f}" == "3,598.7"
assert f"{6 * np.log(len(X)):.1f}" == "60.7" and -drops.max() > 6 * np.log(len(X))


# ================================================================ 6. soft edges ==
with sdt.snippet("soft"):
    resp6 = best["resp"][:, order]
    sure = resp6.max(axis=1)
    print(f"shots at least 99% in one component: {(sure >= 0.99).mean():.1%};"
          f" under 90%: {(sure < 0.9).mean():.1%}; under 60%: {(sure < 0.6).sum():,}")
    split = sure < 0.6
    print("where the split shots are:", shots.BASIC_ZONE[split].value_counts().head(3).to_dict())

    def kmeans(X, k, seed, max_iter=500):
        rng = np.random.default_rng(seed)
        c = X[rng.choice(len(X), k, replace=False)]
        for _ in range(max_iter):
            lab = ((X[:, None] - c[None]) ** 2).sum(axis=2).argmin(axis=1)
            new = np.array([X[lab == j].mean(axis=0) for j in range(k)])
            if np.allclose(new, c):
                break
            c = new
        return lab, c, ((X - c[lab]) ** 2).sum()

    lab_km, c_km, _ = min((kmeans(X, 6, s) for s in range(10)), key=lambda t: t[2])
    table_km = pd.crosstab(lab_km, shots.BASIC_ZONE)
    print(f"k-means, six clusters, ten starts: {table_km.max(axis=1).sum() / len(shots):.1%} in their largest zone")
    for zone in ("Left Corner 3", "Right Corner 3"):
        jk = table_km[zone].idxmax()
        print(f"  the k-means cluster holding the {zone}s: {table_km.loc[jk].sum():,} shots,"
              f" {table_km.loc[jk, zone] / table_km.loc[jk].sum():.1%} of them corner threes")

km_purity = table_km.max(axis=1).sum() / len(shots)
assert f"{km_purity:.1%}" == "67.3%"
assert km_purity < purity - 0.15
assert f"{100 * (purity - km_purity):.1f}" == "18.1"
assert f"{(sure >= 0.99).mean():.1%}" == "50.5%" and f"{(sure < 0.9).mean():.1%}" == "14.1%"
assert int(split.sum()) == 589
vc = shots.BASIC_ZONE[split].value_counts()
assert (vc["In The Paint (Non-RA)"], vc["Above the Break 3"], vc["Mid-Range"]) == (365, 169, 52)
kl, kr = table_km["Left Corner 3"].idxmax(), table_km["Right Corner 3"].idxmax()
assert (int(table_km.loc[kl].sum()), int(table_km.loc[kr].sum())) == (1873, 1894)
assert f"{table_km.loc[kl, 'Left Corner 3'] / table_km.loc[kl].sum():.1%}" == "69.9%"
assert f"{table_km.loc[kr, 'Right Corner 3'] / table_km.loc[kr].sum():.1%}" == "66.2%"
assert table_km.loc[kl, "Left Corner 3"] == 1310 and table_km.loc[kr, "Right Corner 3"] == 1253   # every corner three is there
# the gmm corner strips: how many of their shots are corner threes, and how many corner threes they hold
for j, zone in ((left, "Left Corner 3"), (right, "Right Corner 3")):
    assert table.loc[j, zone] / zc[zone] > 0.97


# ================================================================ 7. the picture ==
def court(ax):
    """Half-court lines in the file's feet: basket at (0, 5.25), sidelines at x = +-25."""
    kw = dict(color="#8A8577", lw=1)
    ax.plot([-25, 25, 25, -25, -25], [0, 0, 47, 47, 0], **kw)
    ax.plot([-22, -22], [0, 5.25 + np.sqrt(23.75 ** 2 - 22 ** 2)], **kw)
    ax.plot([22, 22], [0, 5.25 + np.sqrt(23.75 ** 2 - 22 ** 2)], **kw)
    t = np.degrees(np.arctan2(np.sqrt(23.75 ** 2 - 22 ** 2), 22))
    ax.add_patch(Arc((0, 5.25), 47.5, 47.5, theta1=t, theta2=180 - t, **kw))
    ax.plot([-8, -8, 8, 8], [0, 19, 19, 0], **kw)
    ax.add_patch(Arc((0, 5.25), 1.5, 1.5, **kw))


COLORS = ["#C56A1E", "#2C5E8A", "#6B8E23", "#8E3B46", "#5B4B8A", "#B08D2E"]
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12.6, 5.6), gridspec_kw={"width_ratios": [1.15, 1]})
for j in range(6):
    pts = X[label == j]
    ax1.scatter(pts[:, 0], pts[:, 1], s=2, color=COLORS[j], alpha=0.25, linewidths=0)
for j in range(6):
    vals, vecs = np.linalg.eigh(cov6[j])
    ang = np.degrees(np.arctan2(vecs[1, 1], vecs[0, 1]))
    ax1.add_patch(Ellipse(mean6[j], 4 * np.sqrt(vals[1]), 4 * np.sqrt(vals[0]), angle=ang,
                          fill=False, color=COLORS[j], lw=2))
court(ax1)
ax1.set_xlim(-26, 26)
ax1.set_ylim(-1, 36)
ax1.set_aspect("equal")
ax1.set_xlabel("x (feet across the court)")
ax1.set_ylabel("y (feet from the baseline)")
ax1.set_title(f"Six components, two-sd ellipses: {purity:.1%} of shots in their component's zone",
              fontsize=10)

ks = [s[0] for s in sweep]
ax2.plot(ks, [s[2] / 1000 for s in sweep], marker="o", color="#2C5E8A")
for k, _ll, b, spot in sweep:
    if spot:
        ax2.plot(k, b / 1000, marker="o", color="#C56A1E", ms=9)
ax2.annotate("from k = 8 one component\nsits on the 704 shots logged\nat the middle of the basket",
             xy=(8, sweep[7][2] / 1000), xytext=(4.6, 345), fontsize=9, color="#C56A1E",
             arrowprops=dict(arrowstyle="->", color="#C56A1E"))
ax2.set_xticks(ks)
ax2.set_xlabel("components (k)")
ax2.set_ylabel("BIC, thousands (lower is preferred)")
ax2.set_title("BIC never stops falling", fontsize=10)
fig.tight_layout()
sdt.save_fig(fig, "gmm_court_zones",
             source="NBA_Shots_04_25 shot log, 2023-24 (25,000-shot sample)", asof="June 2026")
plt.close(fig)

print("\nall asserts passed")