import lightgbm as lgb
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import polars as pl
import xgboost as xgb
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import log_loss, roc_auc_score
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, SplineTransformer, StandardScaler

# feature lists and the labelling rule live in the script so notebook and script never drift
from model_baseline import CATEGORICAL, LINEAR, SPLINE
from model_runs import CATS, FEATS, OUTCOMES, expected_runs, frames, label
from evaluation import clustered_log_loss_gain

BLUE, ORANGE, GREY, GREEN = "#2a78d6", "#eb6834", "#8a8a85", "#3f9e6a"
plt.rcParams.update({"figure.dpi": 120, "axes.spines.top": False,
                     "axes.spines.right": False, "axes.grid": True,
                     "grid.alpha": 0.25, "grid.linewidth": 0.6, "font.size": 9})

raw = pl.read_parquet("ball_data.parquet")
df = label(raw)          # legal balls only, runs off the bat -> class index 0..5
print(f"{raw.height:,} deliveries -> {df.height:,} legal  "
      f"({100 * (1 - df.height / raw.height):.1f}% wides/no-balls dropped)")
print(f"{df['match_id'].n_unique():,} matches · {df['year'].min()}-{df['year'].max()}")
print(f"mean runs off the bat: {df['_runs'].mean():.4f} per legal ball")


# Fail early if the generated table no longer matches the notebook's assumptions.
required = {"runs_batter", "is_legal", "match_id", "year", "over", "legal_balls"}
assert not (required - set(raw.columns)), f"missing columns: {sorted(required - set(raw.columns))}"
assert raw["runs_batter"].null_count() == 0
assert set(df["_runs"].unique().to_list()) == set(OUTCOMES)
assert df.select(pl.struct(["match_id", "innings", "over", "ball_in_over"]).n_unique()).item() == df.height

mix = (df.group_by("_runs").agg(pl.len().alias("n"))
         .with_columns((pl.col("n") / df.height).alias("share"),
                       (pl.col("_runs") * pl.col("n")).alias("runs"))
         .with_columns((pl.col("runs") / pl.col("runs").sum()).alias("run_share"))
         .sort("_runs"))
display(mix)

fig, axes = plt.subplots(1, 2, figsize=(9.5, 3.2))
x = np.arange(len(OUTCOMES))
axes[0].bar(x - 0.2, mix["share"], width=0.4, color=BLUE, label="share of balls")
axes[0].bar(x + 0.2, mix["run_share"], width=0.4, color=ORANGE, label="share of runs")
axes[0].set(xticks=x, xlabel="runs off the bat", ylabel="share",
            title="18% of balls, 61% of the runs")
axes[0].set_xticklabels(OUTCOMES)
axes[0].legend(frameon=False, fontsize=8)

by_over = (df.filter(pl.col("over") < 20).group_by("over")
             .agg(pl.col("_runs").mean().alias("runs per ball"),
                  (pl.col("_runs") >= 4).mean().alias("boundary rate"),
                  (pl.col("_runs") == 0).mean().alias("dot rate")).sort("over"))
for col, colour in [("runs per ball", BLUE), ("boundary rate", ORANGE), ("dot rate", GREEN)]:
    axes[1].plot(by_over["over"] + 1, by_over[col], "o-", color=colour, lw=2, ms=4, label=col)
axes[1].set(xlabel="over", title="Powerplay, squeeze, death",
            xticks=range(1, 21, 2))
axes[1].legend(frameon=False, fontsize=8)
plt.tight_layout()

levels, (train, val, test) = frames()
ytr, yva, yte = (d["y"].to_numpy() for d in (train, val, test))

PRIOR = np.bincount(ytr, minlength=len(OUTCOMES)) / len(ytr)   # everything is scored against this
LABELS = list(range(len(OUTCOMES)))
runs_te = np.array(OUTCOMES, dtype=float)[yte]                 # actual runs on each test ball

for name, d, y in [("train", train, ytr), ("val", val, yva), ("test", test, yte)]:
    print(f"{name:6s} {d.height:>7,} balls  {d['year'].min()}-{d['year'].max()}  "
          f"mean runs {np.array(OUTCOMES, float)[y].mean():.4f}")

print("\ntrain prior:  " + "  ".join(f"{v}:{p:.4f}" for v, p in zip(OUTCOMES, PRIOR)))

NULL = np.tile(PRIOR / PRIOR.sum(), (len(yte), 1))
NULL_LL = log_loss(yte, NULL, labels=LABELS)
NULL_RMSE = float(np.sqrt(np.mean((expected_runs(NULL) - runs_te) ** 2)))
print(f"null log loss {NULL_LL:.5f} · null E[runs] RMSE {NULL_RMSE:.4f}")


assert train["year"].max() < val["year"].min() <= val["year"].max() < test["year"].min()
assert set(train["match_id"].unique()).isdisjoint(set(val["match_id"].unique()))
assert set(train["match_id"].unique()).isdisjoint(set(test["match_id"].unique()))

FE = SPLINE + LINEAR + CATEGORICAL
X = lambda d: d.select(FE).to_pandas()

# SimpleImputer runs first: the chase features (rrr, balls_left, ...) are null for every
# first-innings ball, and SplineTransformer cannot take NaN. add_indicator keeps a flag so
# the model still knows the value was missing rather than median.
pre = ColumnTransformer([
    ("spline", Pipeline([
        ("impute", SimpleImputer(strategy="median", add_indicator=True)),
        ("spline", SplineTransformer(n_knots=6, degree=3, include_bias=False)),
        ("scale", StandardScaler()),
    ]), SPLINE),
    ("linear", Pipeline([
        ("impute", SimpleImputer(strategy="median", add_indicator=True)),
        ("scale", StandardScaler()),
    ]), LINEAR),
    ("cat", OneHotEncoder(handle_unknown="ignore", drop="first"), CATEGORICAL),
])

lr = Pipeline([("pre", pre), ("lr", LogisticRegression(max_iter=1000, C=1.0))])
lr.fit(X(train), ytr)

print(f"{len(FE)} features -> {lr.named_steps['pre'].transform(X(val)).shape[1]} columns "
      f"-> {len(OUTCOMES)} class probabilities")

def norm(p):
    """Guard against float drift so log_loss does not complain."""
    return p / p.sum(axis=1, keepdims=True)


def report(name, y, p, runs=None):
    p = norm(p)
    ll = log_loss(y, p, labels=LABELS)
    null_ll = log_loss(y, np.tile(PRIOR, (len(y), 1)), labels=LABELS)
    runs = np.array(OUTCOMES, float)[y] if runs is None else runs
    rmse = np.sqrt(np.mean((expected_runs(p) - runs) ** 2))
    bnd = roc_auc_score((y >= 4).astype(int), p[:, 4] + p[:, 5])
    dot = roc_auc_score((y == 0).astype(int), p[:, 0])
    print(f"{name:6s} logloss {ll:.5f}  vs null {null_ll:.5f} "
          f"({100 * (null_ll - ll) / null_ll:+.2f}%)   E[runs] RMSE {rmse:.4f}   "
          f"boundary AUC {bnd:.4f}   dot AUC {dot:.4f}")


for name, d, y in [("train", train, ytr), ("val", val, yva), ("test", test, yte)]:
    report(name, y, lr.predict_proba(X(d)))

p_lr = norm(lr.predict_proba(X(test)))


def cal_curve(score, actual, bins=10):
    return (pl.DataFrame({"s": score, "a": actual.astype(float)})
              .with_columns(pl.col("s").qcut(bins, labels=[str(i) for i in range(bins)])
                              .alias("b"))
              .group_by("b").agg(pl.len().alias("n"),
                                 pl.col("s").mean().alias("predicted"),
                                 pl.col("a").mean().alias("actual")).sort("b"))


e_lr = cal_curve(expected_runs(p_lr), runs_te)
b_lr = cal_curve(p_lr[:, 4] + p_lr[:, 5], (yte >= 4).astype(int))
display(e_lr)

fig, axes = plt.subplots(1, 2, figsize=(8.6, 4.2))
for ax, c, lo, hi, xl, t in [
        (axes[0], e_lr, 0.8, 2.0, "predicted E[runs]", "Expected runs (deciles)"),
        (axes[1], b_lr, 0.05, 0.32, "predicted P(boundary)", "Boundary probability (deciles)")]:
    ax.plot([lo, hi], [lo, hi], color=GREY, lw=1, ls="--")
    ax.plot(c["predicted"], c["actual"], "o-", color=BLUE, lw=2, ms=6)
    ax.set(xlabel=xl, ylabel="observed", title=t, xlim=(lo, hi), ylim=(lo, hi))
    ax.set_aspect("equal")
plt.tight_layout()

def frame(d, feats=FEATS):
    """Trees take features raw: no splines, no scaling, NaN handled natively."""
    Xd = d.select(feats).to_pandas()
    for c in CATS:
        if c in feats:
            Xd[c] = pd.Categorical(Xd[c], categories=levels[c])
    return Xd


Xtr, Xva, Xte = (frame(d) for d in (train, val, test))

PARAMS = dict(objective="multiclass", num_class=len(OUTCOMES), n_estimators=3000,
              learning_rate=0.03, num_leaves=15, min_child_samples=400, subsample=0.8,
              subsample_freq=1, colsample_bytree=0.7, reg_lambda=10.0, verbose=-1)

lgbm = lgb.LGBMClassifier(**PARAMS, random_state=0)
lgbm.fit(Xtr, ytr, eval_X=Xva, eval_y=yva, eval_metric="multi_logloss",
         callbacks=[lgb.early_stopping(100, verbose=False)])

xgbm = xgb.XGBClassifier(
    objective="multi:softprob", num_class=len(OUTCOMES), n_estimators=3000,
    learning_rate=0.03, max_depth=4, min_child_weight=50, subsample=0.8,
    colsample_bytree=0.7, reg_lambda=10.0, enable_categorical=True, tree_method="hist",
    eval_metric="mlogloss", early_stopping_rounds=100, random_state=0)
xgbm.fit(Xtr, ytr, eval_set=[(Xva, yva)], verbose=False)

print(f"LightGBM stopped at {lgbm.best_iteration_} rounds "
      f"({lgbm.best_iteration_ * len(OUTCOMES)} trees, one per class per round)")
print(f"XGBoost  stopped at {xgbm.best_iteration} rounds")

preds = {
    "null (train prior)": NULL,
    "multinomial + splines": p_lr,
    "lightgbm": norm(lgbm.predict_proba(Xte)),
    "xgboost": norm(xgbm.predict_proba(Xte)),
}
preds["lgb + xgb"] = (preds["lightgbm"] + preds["xgboost"]) / 2

rows = []
for name, p in preds.items():
    p = norm(p)                       # float drift alone is enough to upset log_loss
    ll = log_loss(yte, p, labels=LABELS)
    null_row = name.startswith("null")
    rows.append({"model": name, "logloss": round(ll, 5),
                 "gain_%": round(100 * (NULL_LL - ll) / NULL_LL, 2),
                 "E_runs_RMSE": round(float(np.sqrt(np.mean((expected_runs(p) - runs_te) ** 2))), 4),
                 "boundary_AUC": None if null_row else round(roc_auc_score((yte >= 4).astype(int), p[:, 4] + p[:, 5]), 4),
                 "dot_AUC": None if null_row else round(roc_auc_score((yte == 0).astype(int), p[:, 0]), 4)})
display(pl.DataFrame(rows))

p_lgb = preds["lightgbm"]
rows = []
for i, v in enumerate(OUTCOMES):
    yb = (yte == i).astype(int)
    base = np.full(len(yb), PRIOR[i])
    rows.append({"runs": v, "share": round(float(yb.mean()), 4),
                 "AUC": round(roc_auc_score(yb, p_lgb[:, i]), 4),
                 "logloss": round(log_loss(yb, p_lgb[:, i], labels=[0, 1]), 5),
                 "vs_prior_%": round(100 * (log_loss(yb, base, labels=[0, 1])
                                            - log_loss(yb, p_lgb[:, i], labels=[0, 1]))
                                     / log_loss(yb, base, labels=[0, 1]), 2)})
display(pl.DataFrame(rows))

scored = test.with_columns(pl.Series("e_runs", expected_runs(p_lgb)),
                           pl.Series("p_bnd", p_lgb[:, 4] + p_lgb[:, 5]))


def shape(expr, xlabel, ax, pred, obs, min_n=300):
    g = (scored.with_columns(expr.alias("bucket"))
                .group_by("bucket").agg(pl.len().alias("n"),
                                        pl.col(pred).mean().alias("predicted"),
                                        obs.mean().alias("actual"))
                .filter(pl.col("n") >= min_n).sort("bucket"))
    ax.plot(g["bucket"], g["actual"], "o-", color=BLUE, lw=2, ms=5, label="observed")
    ax.plot(g["bucket"], g["predicted"], "s--", color=ORANGE, lw=2, ms=5, label="model")
    ax.set(xlabel=xlabel)
    ax.legend(frameon=False, fontsize=8)
    return g


fig, axes = plt.subplots(1, 3, figsize=(12.5, 3.4))
shape(pl.col("over") + 1, "over", axes[0], "e_runs", pl.col("_runs"))
axes[0].set(ylabel="runs per ball", title="Runs by over", xticks=range(1, 21, 3))
shape(pl.col("over") + 1, "over", axes[1], "p_bnd", (pl.col("_runs") >= 4).cast(pl.Float64))
axes[1].set(ylabel="P(boundary)", title="Boundary rate by over", xticks=range(1, 21, 3))
shape(pl.col("bat_balls") // 5 * 5, "balls faced by the striker", axes[2], "e_runs", pl.col("_runs"))
axes[2].set(ylabel="runs per ball", title="Getting your eye in")
plt.tight_layout()

gains = sorted(zip(FEATS, lgbm.booster_.feature_importance("gain")), key=lambda t: -t[1])
total = sum(v for _, v in gains)
top = [(k, 100 * v / total) for k, v in gains[:12]][::-1]

fig, ax = plt.subplots(figsize=(6, 3.6))
colours = [ORANGE if k in ("ground", "venue_runs_per_ball", "venue_wicket_rate") else BLUE
           for k, _ in top]
ax.barh([k for k, _ in top], [v for _, v in top], color=colours, height=0.7)
ax.set(xlabel="% of total gain", title="LightGBM importance (orange = venue, see ablation)")
ax.grid(axis="y", visible=False)
plt.tight_layout()

def ablate(drop):
    feats = [f for f in FEATS if f not in drop]
    m = lgb.LGBMClassifier(**PARAMS, random_state=0)
    m.fit(frame(train, feats), ytr, eval_X=frame(val, feats), eval_y=yva,
          eval_metric="multi_logloss", callbacks=[lgb.early_stopping(100, verbose=False)])
    p = norm(m.predict_proba(frame(test, feats)))
    ll = log_loss(yte, p, labels=LABELS)
    return {"dropped": ", ".join(drop) if drop else "(nothing)",
            "logloss": round(ll, 5),
            "gain_%": round(100 * (NULL_LL - ll) / NULL_LL, 2),
            "boundary_AUC": round(roc_auc_score((yte >= 4).astype(int), p[:, 4] + p[:, 5]), 4)}


display(pl.DataFrame([
    ablate([]),
    ablate(["ground"]),
    ablate(["ground", "venue_runs_per_ball", "venue_wicket_rate"]),
    ablate(["h2h_dismissals", "h2h_balls"]),
    ablate(["bat_career_sr", "bat_career_dismissal_rate", "bat_career_balls"]),
]))

rng = np.random.default_rng(0)
ll_base = log_loss(yte, p_lgb, labels=LABELS)
skill = NULL_LL - ll_base


def perm_pct(cols, repeats=3):
    """% of skill lost when `cols` are shuffled together. Returns (mean, sd) over repeats."""
    Xp, losses = Xte.copy(), []
    for _ in range(repeats):
        idx = rng.permutation(len(Xp))      # one index for the whole group, so a collinear set
        for c in cols:                      # stays internally consistent (legal_balls is still
            Xp[c] = Xte[c].values[idx]      # 6 x over) and only its link to the outcome breaks
        losses.append(log_loss(yte, norm(lgbm.predict_proba(Xp)), labels=LABELS))
    for c in cols:
        Xp[c] = Xte[c]
    return 100 * (np.mean(losses) - ll_base) / skill, 100 * np.std(losses) / skill


groups = {f: [f] for f in FEATS}
groups["clock (over+legal_balls+is_powerplay+phase)"] = ["over", "legal_balls", "is_powerplay", "phase"]

rows = []
for name, cols in groups.items():
    pct, sd = perm_pct(cols)
    rows.append({"feature": name, "pct_of_skill": round(pct, 1), "sd": round(sd, 1)})

perm = pl.DataFrame(rows).sort("pct_of_skill", descending=True)
display(perm.head(18))

KEEP = ["is_powerplay", "team_runs", "team_wickets", "legal_balls", "bat_career_sr",
        "balls_since_wicket", "bat_runs", "bat_dot_pct", "bat_balls",
        "bat_career_dismissal_rate", "bowl_career_runs_per_ball", "rrr", "phase",
        "bat_sr", "bat_career_balls"]
print(f"{len(FEATS)} features -> {len(KEEP)}:  {', '.join(KEEP)}\n")

NEW = dict(PARAMS, n_estimators=4000, learning_rate=0.02, min_child_samples=1000,
           reg_lambda=30.0)


def ensemble(feats, params, n_seeds=5):
    """Fit n_seeds models and average their predicted distributions.

    Returns predictions on val as well as test: the calibration step below needs
    a set of predictions the model was not fitted on and that is not the test set.
    """
    a, b, c = (frame(d, feats) for d in (train, val, test))
    pv, pt = [], []
    for s in range(n_seeds):
        m = lgb.LGBMClassifier(**{**params, "random_state": s})
        m.fit(a, ytr, eval_X=b, eval_y=yva, eval_metric="multi_logloss",
              callbacks=[lgb.early_stopping(150, verbose=False)])
        pv.append(norm(m.predict_proba(b)))
        pt.append(norm(m.predict_proba(c)))
    return np.mean(pv, axis=0), np.mean(pt, axis=0)


p_new_va, p_new = ensemble(KEEP, NEW)

rows = []
for name, p in [("null (train prior)", NULL),
                ("multinomial + splines", p_lr),
                ("old: 37 features", p_lgb),
                ("new: 15 features, 5 seeds", p_new)]:
    p = norm(p)
    ll = log_loss(yte, p, labels=LABELS)
    null_row = name.startswith("null")
    rows.append({"model": name, "logloss": round(ll, 5),
                 "gain_%": round(100 * (NULL_LL - ll) / NULL_LL, 2),
                 "E_runs_RMSE": round(float(np.sqrt(np.mean((expected_runs(p) - runs_te) ** 2))), 4),
                 "boundary_AUC": None if null_row else round(roc_auc_score((yte >= 4).astype(int), p[:, 4] + p[:, 5]), 4)})
display(pl.DataFrame(rows))

# Calibration of the final model, both views, against the baseline.
fig, axes = plt.subplots(1, 2, figsize=(8.6, 4.2))
for ax, score, actual, lo, hi, xl, t in [
        (axes[0], lambda p: expected_runs(p), runs_te, 0.8, 2.1,
         "predicted E[runs]", "Expected runs (deciles)"),
        (axes[1], lambda p: p[:, 4] + p[:, 5], (yte >= 4).astype(int), 0.05, 0.33,
         "predicted P(boundary)", "Boundary probability (deciles)")]:
    ax.plot([lo, hi], [lo, hi], color=GREY, lw=1, ls="--")
    for nm, p, colour, mk in [("multinomial + splines", p_lr, BLUE, "o"),
                              ("lgbm, 15 features", p_new, ORANGE, "s")]:
        c = cal_curve(score(p), actual)
        ax.plot(c["predicted"], c["actual"], marker=mk, ls="-", color=colour, lw=2, ms=5, label=nm)
    ax.set(xlabel=xl, ylabel="observed", title=t, xlim=(lo, hi), ylim=(lo, hi))
    ax.set_aspect("equal")
axes[0].legend(frameon=False, fontsize=8, loc="upper left")
plt.tight_layout()

by_year = (pl.concat([train, val, test]).group_by("year")
             .agg(pl.col("_runs").mean().alias("actual")).sort("year"))
e_new = expected_runs(p_new)
print(f"train mean {np.array(OUTCOMES, float)[ytr].mean():.4f}  ->  "
      f"test mean {runs_te.mean():.4f}  (+{100 * (runs_te.mean() / np.array(OUTCOMES, float)[ytr].mean() - 1):.1f}%)")
print(f"model mean prediction on test: {e_new.mean():.4f}  "
      f"({e_new.mean() - runs_te.mean():+.4f} per ball)")

fig, ax = plt.subplots(figsize=(6.5, 3))
ax.plot(by_year["year"], by_year["actual"], "o-", color=BLUE, lw=2, ms=4, label="observed")
ax.axvspan(2024.5, 2026.5, color=ORANGE, alpha=0.10)
ax.scatter(test.select(pl.col("year")).to_series().unique().sort(),
           [pl.DataFrame({"y": test["year"], "e": e_new}).filter(pl.col("y") == yr)["e"].mean()
            for yr in sorted(test["year"].unique().to_list())],
           color=ORANGE, marker="s", zorder=3, label="model, test seasons")
ax.set(xlabel="season", ylabel="runs per ball", title="Scoring drifted up; the model did not")
ax.legend(frameon=False, fontsize=8)
plt.tight_layout()

def prior_shift(p_fit, y_fit):
    """Per-class multiplicative correction: observed frequency / mean predicted probability."""
    obs = np.bincount(y_fit, minlength=len(OUTCOMES)) / len(y_fit)
    return obs / p_fit.mean(axis=0)


def temper(p, T):
    z = np.log(np.clip(p, 1e-12, None)) / T
    z -= z.max(axis=1, keepdims=True)
    return norm(np.exp(z))


W = prior_shift(p_new_va, yva)
grid = np.linspace(0.6, 1.6, 101)
T = float(grid[int(np.argmin([log_loss(yva, temper(p_new_va, t), labels=LABELS) for t in grid]))])

print("per-class weights fitted on 2024:")
for v, w in zip(OUTCOMES, W):
    print(f"  {v}: {w:.3f}")
print(f"\nfitted temperature: {T:.2f}")

p_cal = norm(p_new * W)

rows = []
for name, p in [("null (train prior)", NULL),
                ("15 features, uncalibrated", p_new),
                (f"+ temperature (T={T:.2f})", temper(p_new, T)),
                ("+ prior shift", p_cal)]:
    p = norm(p)
    ll = log_loss(yte, p, labels=LABELS)
    null_row = name.startswith("null")
    rows.append({"model": name, "logloss": round(ll, 5),
                 "gain_%": round(100 * (NULL_LL - ll) / NULL_LL, 2),
                 "mean_E_runs": round(float(expected_runs(p).mean()), 4),
                 "E_runs_RMSE": round(float(np.sqrt(np.mean((expected_runs(p) - runs_te) ** 2))), 4),
                 "boundary_AUC": None if null_row else round(roc_auc_score((yte >= 4).astype(int), p[:, 4] + p[:, 5]), 4)})
display(pl.DataFrame(rows))
print(f"actual mean runs per ball on test: {runs_te.mean():.4f}")

fig, axes = plt.subplots(1, 2, figsize=(8.6, 4.2))
for ax, score, actual, lo, hi, xl, t in [
        (axes[0], lambda p: expected_runs(p), runs_te, 0.8, 2.3,
         "predicted E[runs]", "Expected runs (deciles)"),
        (axes[1], lambda p: p[:, 4] + p[:, 5], (yte >= 4).astype(int), 0.05, 0.36,
         "predicted P(boundary)", "Boundary probability (deciles)")]:
    ax.plot([lo, hi], [lo, hi], color=GREY, lw=1, ls="--")
    for nm, p, colour, mk in [("uncalibrated", p_new, ORANGE, "s"),
                              ("prior shift", p_cal, GREEN, "D")]:
        cc = cal_curve(score(p), actual)
        ax.plot(cc["predicted"], cc["actual"], marker=mk, ls="-", color=colour, lw=2, ms=5, label=nm)
    ax.set(xlabel=xl, ylabel="observed", title=t, xlim=(lo, hi), ylim=(lo, hi))
    ax.set_aspect("equal")
axes[0].legend(frameon=False, fontsize=8, loc="upper left")
plt.tight_layout()

mid = test.filter(pl.col("innings") == 2)["match_id"][0]
inn = (test.with_columns(pl.Series("e_runs", expected_runs(p_cal)),
                         pl.Series("p_bnd", p_cal[:, 4] + p_cal[:, 5]))
           .filter((pl.col("match_id") == mid) & (pl.col("innings") == 2))
           .with_row_index("ball"))

fig, axes = plt.subplots(2, 1, figsize=(9.5, 4.6), sharex=True,
                         gridspec_kw={"height_ratios": [2, 1]})
axes[0].plot(inn["ball"], inn["e_runs"], color=BLUE, lw=1.6, label="E[runs]")
ax2 = axes[0].twinx()
ax2.plot(inn["ball"], inn["p_bnd"], color=ORANGE, lw=1.2, alpha=0.8, label="P(boundary)")
ax2.set_ylabel("P(boundary)", color=ORANGE)
ax2.grid(False)
axes[0].set(ylabel="E[runs]", title=f"{inn['batting_team'][0]} chasing "
            f"{inn['target_runs'][0]} v {inn['bowling_team'][0]} ({inn['date'][0]})")
axes[1].bar(inn["ball"], inn["_runs"], color=GREY, width=0.9)
axes[1].set(xlabel="ball of the innings", ylabel="runs", yticks=[0, 2, 4, 6])
plt.tight_layout()

style = df                      # already legal-only, with _runs, from the first cell

fig, axes = plt.subplots(1, 3, figsize=(12.5, 3.4))

# --- when each type bowls, which is what the model has been inferring ---
mix = (style.filter(pl.col("over") < 20).group_by("over")
            .agg((pl.col("bowler_type") == "spin").mean().alias("spin")).sort("over"))
axes[0].bar(mix["over"] + 1, mix["spin"], color=BLUE, width=0.72)
axes[0].set(xlabel="over", ylabel="share of balls bowled by spin",
            title="Spin bowls the middle", xticks=range(1, 21, 3))

# --- runs by type and phase ---
g = (style.group_by(["bowler_type", "is_powerplay"])
          .agg(pl.col("_runs").mean().alias("rpb"),
               (pl.col("_runs") >= 4).mean().alias("bnd")).sort("bowler_type", "is_powerplay"))
x = np.arange(2)
for i, (t, colour) in enumerate([("pace", ORANGE), ("spin", GREEN)]):
    v = g.filter(pl.col("bowler_type") == t).sort("is_powerplay")
    axes[1].bar(x + (i - 0.5) * 0.38, v["rpb"], width=0.36, color=colour, label=t)
axes[1].set(xticks=x, ylabel="runs per ball", title="The gap is a middle-overs gap")
axes[1].set_xticklabels(["overs 7-20", "powerplay"])
axes[1].legend(frameon=False, fontsize=8)

# --- the matchup ---
m = (style.filter(pl.col("bowler_type") == "spin")
          .drop_nulls("turn_into_batter").group_by("turn_into_batter")
          .agg(pl.len().alias("n"), pl.col("_runs").mean().alias("rpb"),
               (pl.col("_runs") >= 4).mean().alias("bnd")).sort("turn_into_batter"))
axes[2].bar(["turns away", "turns in"], m["rpb"], color=[GREY, BLUE], width=0.55)
for i, (r, n) in enumerate(zip(m["rpb"], m["n"])):
    axes[2].annotate(f"{r:.3f}\n({n:,} balls)", (i, r), ha="center", va="bottom",
                     xytext=(0, 2), textcoords="offset points", fontsize=8)
axes[2].set(ylabel="runs per ball", title="Spin: which way it turns", ylim=(0, 1.55))
plt.tight_layout()

display(style.group_by("bowler_type").agg(
    pl.len().alias("balls"), pl.col("_runs").mean().round(4).alias("runs_per_ball"),
    (pl.col("_runs") >= 4).mean().round(4).alias("boundary_rate"),
    (pl.col("_runs") == 0).mean().round(4).alias("dot_rate"),
    pl.col("wicket").mean().round(4).alias("wicket_rate")).sort("bowler_type"))