"""
Tutorial 87 - Predict a scoreline with Poisson: the 2026 World Cup final as a grid.

Reads the bundled wc2026_results.csv - the complete 2026 World Cup, all 104
matches from the group stage through the final - so it runs offline. The build
is the classic Poisson scoreline model, from scratch in numpy and pandas with
the pmf written out by hand (no scipy): a league-average goal rate first, the
honest check of that single-rate model against the real goals-per-team
distribution (it fails, narrowly and instructively - var/mean is 1.31 because
teams differ), then per-team attack and defence multipliers, and finally the
payoff grid: the score matrix for Spain vs Argentina, P(Spain scores h) times
P(Argentina scores a) for every scoreline. The most likely cell is 1-0 Spain
at about 30%% - which is exactly how the real final ended. We finish with the
model's known blind spot, measured rather than asserted: independence
under-predicts draws (15.1 expected in the group stage, 20 observed), the
Dixon-Coles correction's reason to exist. Deterministic end to end - every
number recomputes identically on every run.

Run:  python downloads/87_predict_a_scoreline_with_poisson.py
"""
import math
import os

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

import sdt_common as sdt

sdt.init("predict-a-scoreline-with-poisson")
HERE = os.path.dirname(os.path.abspath(__file__))

df = pd.read_csv(os.path.join(HERE, "wc2026_results.csv"))
assert len(df) == 104, f"expected the complete tournament, got {len(df)}"


def pois(k, lam):
    """P(X = k) for X ~ Poisson(lam) - the pmf, written out instead of imported."""
    return math.exp(-lam) * lam ** k / math.factorial(k)


# --- one row per TEAM per match: 104 matches -> 208 scoring opportunities ----------
long = pd.concat([
    pd.DataFrame({"stage": df["stage"], "team": df["home"],
                  "scored": df["home_goals"], "conceded": df["away_goals"]}),
    pd.DataFrame({"stage": df["stage"], "team": df["away"],
                  "scored": df["away_goals"], "conceded": df["home_goals"]}),
], ignore_index=True)
assert len(long) == 208

lam0 = long["scored"].mean()
var0 = long["scored"].var(ddof=1)
assert 1.40 < lam0 < 1.55, f"lam0 moved: {lam0}"
assert 1.25 < var0 / lam0 < 1.40, f"dispersion moved: {var0 / lam0}"

with sdt.snippet("lambda"):
    print(f"{len(df)} matches -> {len(long)} team-matches\n")
    print(f"  goals per team per match (lambda)   {lam0:.3f}")
    print(f"  variance of goals per team-match    {var0:.3f}")
    print(f"  variance / mean                     {var0 / lam0:.2f}")
    print()
    print("a Poisson variable's variance EQUALS its mean, so that ratio should")
    print("be ~1.00. we got 1.31 - our first hint that one shared rate is too")
    print("simple. hold the thought; first, how close does one rate get?")

# --- observed vs Poisson(lam0): the honest single-rate check ----------------------
obs = long["scored"].value_counts().reindex(range(8), fill_value=0)
expected = np.array([len(long) * pois(k, lam0) for k in range(8)])

# chi-square on 6 bins (0,1,2,3,4,5+), one estimated parameter -> df = 4
obs_b = list(obs[:5]) + [int(obs[5:].sum())]
exp_b = list(expected[:5]) + [len(long) - expected[:5].sum()]
chi2 = sum((o - e) ** 2 / e for o, e in zip(obs_b, exp_b))
CRIT = 9.488  # chi-square critical value, df=4, alpha=0.05
assert 9.0 < chi2 < 10.0, f"chi-square moved: {chi2}"
assert chi2 > CRIT, "the borderline rejection is part of the story"

with sdt.snippet("fit"):
    print("goals scored by one team in one match - observed vs Poisson(1.481):\n")
    print("   goals   observed   Poisson expects")
    for k in range(8):
        print(f"     {k}       {obs[k]:>4}         {expected[k]:>6.1f}")
    print()
    print("close - the shape is unmistakably Poisson-ish - but look at the")
    print("misses: too many 0s (55 vs 47), too few 2s (39 vs 52), too many")
    print("blowout tallies of 5+. pooling bins 0-4 and 5+:")
    print(f"\n  chi-square = {chi2:.2f}  vs the df=4, 5% bar of {CRIT}")
    print()
    print("a hair over the line. the single-rate model fails, and for a")
    print("physical reason: 48 different teams don't share one lambda. mixing")
    print("strong and weak attacks fattens both tails - that's the var/mean of")
    print("1.31 showing up again. the fix is one rate per team.")

# --- figure 1: the distribution check ---------------------------------------------
color = sdt.sport_color("soccer")
fig, ax = plt.subplots(figsize=(9.6, 5.2))
ks = np.arange(8)
ax.bar(ks, obs.to_numpy(), color=color, edgecolor="#16243f", linewidth=0.8,
       zorder=3, label="observed (208 team-matches)")
ax.plot(ks, expected, color="#20242B", lw=1.6, marker="o", markersize=6,
        zorder=4, label=f"Poisson, one shared rate ($\\lambda$ = {lam0:.2f})")
for k in (0, 2):
    ax.annotate("", xy=(k, expected[k]), xytext=(k, obs[k]),
                arrowprops=dict(arrowstyle="->", color="#bd4b34", lw=1.6), zorder=5)
ax.annotate("too many 0s, too few 2s:\nteams differ", xy=(2.55, 44), fontsize=9.5,
            color="#bd4b34", ha="left")
ax.set_xlabel("goals scored by one team in one match")
ax.set_ylabel("team-matches")
ax.set_title("One Poisson rate almost fits the whole World Cup - almost")
ax.legend(frameon=False)
sdt.save_fig(fig, "goals_fit",
             source="Bundled complete 2026 World Cup results (ESPN public data)",
             asof="July 2026")

# --- per-team attack and defence multipliers --------------------------------------
rates = (
    long.groupby("team")
        .agg(games=("scored", "size"), scored=("scored", "mean"),
             conceded=("conceded", "mean"))
)
rates["attack"] = rates["scored"] / lam0     # >1 = scores more than average
rates["defence"] = rates["conceded"] / lam0  # <1 = concedes less than average

with sdt.snippet("rates"):
    print("attack = goals scored per match / 1.481   (1 = tournament average)")
    print("defence = goals conceded per match / 1.481 (lower = stingier)\n")
    print("sharpest attacks:")
    print(rates.sort_values("attack", ascending=False).head(4).round(3).to_string())
    print("\nstingiest defences:")
    print(rates.sort_values("defence").head(4).round(3).to_string())
    print()
    print("Spain conceded ONCE in eight games - defence 0.084. and note the")
    print("games column: finalists played 8, group-stage exits played 3.")
    print("Panama scored 0 goals in 3 games, so their attack multiplier is")
    print("0.000 - the model now claims Panama can literally never score.")
    print("three games of evidence, taken at face value. remember that.")

assert rates.loc["Panama", "attack"] == 0.0
assert rates.loc["Spain", "defence"] < 0.1

# --- the score matrix for the final -----------------------------------------------
HOME, AWAY = "Spain", "Argentina"
lam_h = lam0 * rates.loc[HOME, "attack"] * rates.loc[AWAY, "defence"]
lam_a = lam0 * rates.loc[AWAY, "attack"] * rates.loc[HOME, "defence"]

K = 7  # model scorelines 0..6; beyond that the probabilities are dust
p_h = np.array([pois(k, lam_h) for k in range(K)])
p_a = np.array([pois(k, lam_a) for k in range(K)])
M = np.outer(p_h, p_a)  # M[h, a] = P(Spain h) * P(Argentina a)

hi, ai = np.unravel_index(M.argmax(), M.shape)
p_win = np.tril(M, -1).sum()   # Spain scores more: below the diagonal
p_draw = np.trace(M)
p_loss = np.triu(M, 1).sum()
assert (hi, ai) == (1, 0), f"most likely scoreline moved: {hi}-{ai}"
assert 0.28 < M[1, 0] < 0.32
assert 0.60 < p_win < 0.65 and p_loss < 0.10
assert M.sum() > 0.995  # the 7x7 grid captures essentially all probability

with sdt.snippet("final"):
    print(f"the final, {HOME} vs {AWAY}:\n")
    print(f"  lambda_{HOME}     = 1.481 x {rates.loc[HOME, 'attack']:.3f} x "
          f"{rates.loc[AWAY, 'defence']:.3f} = {lam_h:.2f} expected goals")
    print(f"  lambda_{AWAY} = 1.481 x {rates.loc[AWAY, 'attack']:.3f} x "
          f"{rates.loc[HOME, 'defence']:.3f} = {lam_a:.2f} expected goals\n")
    flat = sorted(((M[i, j], i, j) for i in range(K) for j in range(K)), reverse=True)
    print("most likely scorelines:")
    for p, i, j in flat[:5]:
        print(f"    {HOME} {i}-{j}  {p:.1%}")
    print()
    print(f"  {HOME} win  {p_win:.1%}   draw  {p_draw:.1%}   {AWAY} win  {p_loss:.1%}")
    print(f"  (the {K}x{K} grid holds {M.sum():.1%} of all probability)\n")
    print("the model's top cell is 1-0 Spain. the real final: Spain 1-0.")

# --- figure 2: the score matrix as a heatmap --------------------------------------
fig, ax = plt.subplots(figsize=(8.6, 7.0))
im = ax.imshow(M, cmap="Greens", vmin=0, vmax=M.max(), aspect="equal")
ax.grid(False)
for i in range(K):
    for j in range(K):
        dark = M[i, j] > 0.55 * M.max()
        ax.text(j, i, f"{M[i, j] * 100:.1f}", ha="center", va="center",
                fontsize=9.5, color="#FBF7EE" if dark else "#20242B",
                fontweight="bold" if (i, j) == (hi, ai) else "normal")
ax.add_patch(plt.Rectangle((ai - 0.5, hi - 0.5), 1, 1, fill=False,
                           edgecolor="#bd4b34", lw=2.4, zorder=5))
ax.annotate("the actual result", xy=(ai + 0.45, hi), xytext=(1.7, 0.25),
            fontsize=10, color="#bd4b34", fontweight="bold",
            arrowprops=dict(arrowstyle="->", color="#bd4b34", lw=1.4))
ax.set_xticks(range(K))
ax.set_yticks(range(K))
ax.set_xlabel(f"{AWAY} goals")
ax.set_ylabel(f"{HOME} goals")
ax.set_title("Every scoreline's probability: Spain vs Argentina, the 2026 final\n"
             f"(cell = % chance; Spain win {p_win:.0%}, draw {p_draw:.0%}, "
             f"Argentina win {p_loss:.0%})")
cb = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.03)
cb.set_label("probability of this exact scoreline", fontsize=9)
sdt.save_fig(fig, "score_matrix",
             source="Poisson model fit on the bundled complete 2026 World Cup "
                    "results (ESPN public data)",
             asof="July 2026")

# --- where independence breaks: the draw ledger -----------------------------------
grp = df[df["stage"] == "group"].copy()
drawn = int((grp["home_goals"] == grp["away_goals"]).sum())

exp_draws = 0.0
for _, m in grp.iterrows():
    lh = lam0 * rates.loc[m["home"], "attack"] * rates.loc[m["away"], "defence"]
    la = lam0 * rates.loc[m["away"], "attack"] * rates.loc[m["home"], "defence"]
    ph = np.array([pois(k, lh) for k in range(K)])
    pa = np.array([pois(k, la) for k in range(K)])
    exp_draws += float(np.outer(ph, pa).trace())

assert drawn == 20 and 14.5 < exp_draws < 15.5

with sdt.snippet("draws"):
    print("run the model over all 72 group games (the only stage that can end")
    print("level) and add up each game's draw probability:\n")
    print(f"  draws the model expects   {exp_draws:.1f}")
    print(f"  draws that happened       {drawn}\n")
    print("the shortfall is systematic, not bad luck: multiplying the two")
    print("Poissons assumes the sides score INDEPENDENTLY, but real teams")
    print("respond to the scoreboard - level games tighten up, 1-1 protects")
    print("itself. every serious scoreline model patches exactly this cell")
    print("(Dixon & Coles, 1997); the patch is the challenge below.")

# --- the pre-final refit quoted in the troubleshooting section --------------------
# (fit on the 103 matches BEFORE the final, then genuinely forecast it)
pre = df[df["stage"] != "final"]
plong = pd.concat([
    pd.DataFrame({"team": pre["home"], "scored": pre["home_goals"],
                  "conceded": pre["away_goals"]}),
    pd.DataFrame({"team": pre["away"], "scored": pre["away_goals"],
                  "conceded": pre["home_goals"]}),
], ignore_index=True)
plam0 = plong["scored"].mean()
pr = plong.groupby("team").agg(scored=("scored", "mean"), conceded=("conceded", "mean"))
plh = plam0 * (pr.loc[HOME, "scored"] / plam0) * (pr.loc[AWAY, "conceded"] / plam0)
pla = plam0 * (pr.loc[AWAY, "scored"] / plam0) * (pr.loc[HOME, "conceded"] / plam0)
pM = np.outer([pois(k, plh) for k in range(K)], [pois(k, pla) for k in range(K)])
assert np.unravel_index(pM.argmax(), pM.shape) == (1, 0), "held-out top cell moved"
assert 0.26 < pM[1, 0] < 0.29, f"held-out 1-0 prob moved: {pM[1, 0]}"
assert 0.60 < np.tril(pM, -1).sum() < 0.65, "held-out Spain win prob moved"

print("done")