bayesian/simulator.py

"""Monte Carlo tournament simulator for the 2026 FIFA World Cup.

Simulates the full tournament N times:
  - Group stage: round-robin, played matches use actual scores
  - Qualifier rule: top 2 per group + 8 best 3rd-place teams → 32 teams
  - Knockout: single-elimination with 50/50 penalty shootout on draws

Returns championship probability for every team.
"""
from __future__ import annotations

from collections import Counter, defaultdict

import numpy as np

from data.live import get_2026_fixtures
from bayesian.model import _to_hist_name

N_SIMS = 5_000


# ── lambda matrix pre-computation ────────────────────────────────────────────

def build_lambda_matrix(params: dict, team_idx: dict[str, int]) -> tuple[np.ndarray, np.ndarray]:
    """Pre-compute expected-goals matrices for all team pairs.

    lh_mat[i, j] = E[home goals]  when team-i is home, team-j is away
    la_mat[i, j] = E[away goals]  when team-i is home, team-j is away
    """
    n = len(team_idx)
    atk = params["attack"]  # shape (n,)
    dfs = params["defense"]  # shape (n,)
    mu = params["mu"]
    gamma = params["gamma"]

    lh_mat = np.exp(np.clip(mu + gamma + atk[:, None] - dfs[None, :], -5, 5))
    la_mat = np.exp(np.clip(mu + atk[None, :] - dfs[:, None], -5, 5))
    return lh_mat, la_mat


# ── fixture parsing ───────────────────────────────────────────────────────────

def _parse_score(match: dict) -> tuple[int, int] | None:
    score = match.get("score")
    if not score:
        return None
    ft = score.get("ft", [])
    return (int(ft[0]), int(ft[1])) if len(ft) >= 2 else None


def build_group_data(team_idx: dict[str, int]) -> dict:
    """Parse 2026 group-stage fixtures into a structured dict.

    Returns:
        {
          "Group A": {
            "teams": [hist_name, ...],      # 4 teams in hist-name space
            "fixtures": [
              (home_idx, away_idx, actual_hg_or_None, actual_ag_or_None),
              ...
            ]
          },
          ...
        }
    """
    all_fixtures = get_2026_fixtures()
    group_fixtures: dict[str, list] = defaultdict(list)
    group_teams: dict[str, set] = defaultdict(set)

    for m in all_fixtures:
        grp = m.get("group", "")
        if not grp:
            continue
        h_raw = _to_hist_name(m["team1"])
        a_raw = _to_hist_name(m["team2"])
        group_teams[grp].add(h_raw)
        group_teams[grp].add(a_raw)
        actual = _parse_score(m)
        h_idx = team_idx.get(h_raw)
        a_idx = team_idx.get(a_raw)
        if h_idx is None or a_idx is None:
            continue
        group_fixtures[grp].append((
            h_idx, a_idx,
            actual[0] if actual else None,
            actual[1] if actual else None,
        ))

    return {
        grp: {
            "teams": [team_idx[t] for t in sorted(group_teams[grp]) if t in team_idx],
            "fixtures": group_fixtures[grp],
        }
        for grp in sorted(group_fixtures.keys())
    }


# ── group-stage simulation ────────────────────────────────────────────────────

def _group_standings(team_indices: list[int], results: list[tuple]) -> list[tuple]:
    """Return teams sorted by (points, goal_diff, goals_for) descending.

    results: list of (home_idx, away_idx, hg, ag)
    Returns: list of (team_idx, pts, gd, gf)
    """
    pts: dict[int, int] = {t: 0 for t in team_indices}
    gd: dict[int, int] = {t: 0 for t in team_indices}
    gf: dict[int, int] = {t: 0 for t in team_indices}

    for h, a, hg, ag in results:
        gf[h] += hg
        gf[a] += ag
        gd[h] += hg - ag
        gd[a] += ag - hg
        if hg > ag:
            pts[h] += 3
        elif ag > hg:
            pts[a] += 3
        else:
            pts[h] += 1
            pts[a] += 1

    ranked = sorted(team_indices, key=lambda t: (pts[t], gd[t], gf[t]), reverse=True)
    return [(t, pts[t], gd[t], gf[t]) for t in ranked]


# ── knockout simulation ───────────────────────────────────────────────────────

def _ko_match(h: int, a: int, lh_mat: np.ndarray, la_mat: np.ndarray, rng: np.random.Generator) -> int:
    """Simulate a knockout match. Draws go to 50/50 penalty shootout."""
    hg = rng.poisson(lh_mat[h, a])
    ag = rng.poisson(la_mat[h, a])
    if hg > ag:
        return h
    if ag > hg:
        return a
    return h if rng.random() < 0.5 else a


def _simulate_knockout_bracket(teams: list[int], lh_mat: np.ndarray, la_mat: np.ndarray, rng: np.random.Generator) -> int:
    """Single-elimination bracket starting from `teams` (must be power of 2 in length).

    Teams are paired sequentially: (0 vs 1), (2 vs 3), …
    """
    round_teams = list(teams)
    while len(round_teams) > 1:
        next_round = []
        for i in range(0, len(round_teams), 2):
            winner = _ko_match(round_teams[i], round_teams[i + 1], lh_mat, la_mat, rng)
            next_round.append(winner)
        round_teams = next_round
    return round_teams[0]


# ── full tournament simulation ────────────────────────────────────────────────

def simulate_tournament(
    params: dict,
    team_idx: dict[str, int],
    n_sims: int = N_SIMS,
    seed: int = 42,
) -> dict[str, float]:
    """Run n_sims Monte Carlo simulations of the full 2026 WC.

    Returns dict of {team_hist_name: championship_probability}.
    """
    rng = np.random.default_rng(seed)
    lh_mat, la_mat = build_lambda_matrix(params, team_idx)
    group_data = build_group_data(team_idx)

    if not group_data:
        return {}

    idx_to_team = {v: k for k, v in team_idx.items()}
    winner_counts: Counter = Counter()

    for _ in range(n_sims):
        all_third: list[tuple] = []
        auto_qualifiers: list[int] = []

        # ---- group stage ----
        for grp_name, grp in group_data.items():
            team_indices = grp["teams"]
            results = []
            for h_idx, a_idx, actual_hg, actual_ag in grp["fixtures"]:
                if actual_hg is not None:
                    results.append((h_idx, a_idx, actual_hg, actual_ag))
                else:
                    hg = int(rng.poisson(lh_mat[h_idx, a_idx]))
                    ag = int(rng.poisson(la_mat[h_idx, a_idx]))
                    results.append((h_idx, a_idx, hg, ag))

            standings = _group_standings(team_indices, results)
            auto_qualifiers.append(standings[0][0])  # 1st place
            auto_qualifiers.append(standings[1][0])  # 2nd place
            all_third.append(standings[2])            # 3rd place entry

        # ---- 8 best 3rd-place teams ----
        best_third = sorted(all_third, key=lambda x: (x[1], x[2], x[3]), reverse=True)[:8]
        r32_teams = auto_qualifiers + [entry[0] for entry in best_third]

        # ---- knockout bracket ----
        # Shuffle so bracket pairings vary across simulations (no fixed seeding bias)
        rng.shuffle(r32_teams)
        champion_idx = _simulate_knockout_bracket(r32_teams, lh_mat, la_mat, rng)
        winner_counts[idx_to_team[champion_idx]] += 1

    total = sum(winner_counts.values())
    return {team: count / total for team, count in winner_counts.most_common()}