bayesian/model.py

"""Dixon-Coles Bayesian Poisson model for World Cup match prediction.

Data sources used for fitting:
  1. Historical WC matches 1930-2022 (jfjelstul/worldcup dataset)
  2. Recent international results 2022-present (martj42/international_results)

Priors:
  - FIFA World Rankings (June 2025) provide informative prior means for
    each team's total strength (α+δ), split evenly between attack and defense.
  - L2 regularisation (λ=1.0) around those prior means acts as a Gaussian
    prior: αᵢ ~ N(prior_atk_i, 1),  δᵢ ~ N(prior_dfs_i, 1).

Uncertainty is the Laplace approximation (inverse Hessian diagonal at MAP).
"""
import io
import math
import time
import unicodedata
from pathlib import Path
from threading import Lock
from typing import Any

import numpy as np
import pandas as pd
import requests
from scipy.optimize import minimize
from scipy.stats import poisson

from data.loader import get_matches
from data.live import get_all_named_matches, normalize_team_name
from bayesian.rankings import get_prior_means

_DECAY_XI = 0.002      # per-week temporal decay
_MAX_GOALS = 8         # score matrix truncation
_RIDGE_LAMBDA = 1.0    # L2 regularisation strength
_DC_RHO = -0.1         # Dixon-Coles low-score correction

# Recent data settings
_RECENT_URL = (
    "https://raw.githubusercontent.com/martj42/international_results"
    "/master/results.csv"
)
_RECENT_CACHE = Path(__file__).parent.parent / "data" / "2026" / "recent_results.csv"
_RECENT_FROM = "2022-01-01"     # only use post-2022 WC cycle data
_RECENT_CACHE_LOCK = Lock()
_RECENT_CACHE_TTL = 86400       # refresh local cache once per day
_FRIENDLY_WEIGHT = 0.4          # down-weight friendlies vs competitive matches

_model_cache: dict | None = None


# ── name normalisation ────────────────────────────────────────────────────────

def _norm(name: str) -> str:
    nfkd = unicodedata.normalize("NFKD", name)
    return "".join(c for c in nfkd if not unicodedata.combining(c)).lower().strip()


# Maps openfootball / martj42 / common names → jfjelstul historical names
_HIST_NAME_MAP: dict[str, str] = {
    "usa": "United States",
    "united states of america": "United States",
    "bosnia & herzegovina": "Bosnia and Herzegovina",
    "bosnia-herzegovina": "Bosnia and Herzegovina",
    "bosnia and herzegovina": "Bosnia and Herzegovina",
    "ir iran": "Iran",
    "korea republic": "South Korea",
    "republic of korea": "South Korea",
    "dpr korea": "North Korea",
    "trinidad & tobago": "Trinidad and Tobago",
    "trinidad and tobago": "Trinidad and Tobago",
    "côte d'ivoire": "Ivory Coast",
    "cote d'ivoire": "Ivory Coast",
    "ivory coast": "Ivory Coast",
    "china pr": "China",
    "chinese taipei": "Taiwan",
    "north macedonia": "Macedonia",
    "republic of ireland": "Ireland",
}


def _to_hist_name(name: str) -> str:
    """Map any external team name to the name used in this model's team index."""
    direct = normalize_team_name(name)
    normed = _norm(direct)
    return _HIST_NAME_MAP.get(normed, direct)


# ── recent international data ─────────────────────────────────────────────────

def _is_cache_fresh() -> bool:
    if not _RECENT_CACHE.exists():
        return False
    age = time.time() - _RECENT_CACHE.stat().st_mtime
    return age < _RECENT_CACHE_TTL


def fetch_recent_international_data() -> pd.DataFrame:
    """Download martj42 international results CSV and return post-2022 rows.

    Caches locally to avoid re-downloading on every model fit.
    Returns an empty DataFrame on any failure so the model degrades gracefully.
    """
    with _RECENT_CACHE_LOCK:
        if not _is_cache_fresh():
            try:
                resp = requests.get(_RECENT_URL, timeout=30)
                resp.raise_for_status()
                _RECENT_CACHE.parent.mkdir(parents=True, exist_ok=True)
                _RECENT_CACHE.write_bytes(resp.content)
            except Exception:
                if not _RECENT_CACHE.exists():
                    return pd.DataFrame()

        try:
            df = pd.read_csv(_RECENT_CACHE, low_memory=False)
        except Exception:
            return pd.DataFrame()

    required = {"date", "home_team", "away_team", "home_score", "away_score", "tournament"}
    if not required.issubset(df.columns):
        return pd.DataFrame()

    df["date"] = pd.to_datetime(df["date"], errors="coerce")
    df = df.dropna(subset=["date", "home_score", "away_score"])
    df = df[df["date"] >= _RECENT_FROM].copy()
    df["home_score"] = df["home_score"].astype(int)
    df["away_score"] = df["away_score"].astype(int)

    # Normalise team names to the same space as the historical WC data
    df["home_team"] = df["home_team"].apply(_to_hist_name)
    df["away_team"] = df["away_team"].apply(_to_hist_name)

    return df


def _recent_match_weight(tournament: str) -> float:
    """Weight competitive matches higher than friendlies."""
    t = str(tournament).lower()
    if "friendly" in t:
        return _FRIENDLY_WEIGHT
    return 1.0


# ── historical WC data ────────────────────────────────────────────────────────

def _get_wc2026_team_names() -> set[str]:
    """Return the set of normalised team names playing in the 2026 WC."""
    try:
        matches = get_all_named_matches()
        names: set[str] = set()
        for m in matches:
            names.add(_to_hist_name(m["team1"]))
            names.add(_to_hist_name(m["team2"]))
        return names
    except Exception:
        return set()


def load_training_data() -> dict:
    """Load and merge historical WC + recent international data for model fitting.

    Recent data is filtered to matches where at least one team is a 2026 WC
    participant, keeping the parameter space tractable (~100-150 teams vs 271).
    """
    # --- Historical WC matches (1930-2022) ---
    hist = get_matches()
    hist = hist[hist["tournament_name"].str.contains("Men", na=False)].copy()
    hist = hist.dropna(subset=["home_team_score", "away_team_score"])
    hist["home_team_score"] = hist["home_team_score"].astype(int)
    hist["away_team_score"] = hist["away_team_score"].astype(int)
    hist["match_date"] = pd.to_datetime(hist["match_date"], errors="coerce")
    hist = hist.dropna(subset=["match_date"])

    rows_hist = pd.DataFrame({
        "date": hist["match_date"],
        "home_team": hist["home_team_name"],
        "away_team": hist["away_team_name"],
        "home_score": hist["home_team_score"],
        "away_score": hist["away_team_score"],
        "base_weight": 1.0,
    })

    # --- Recent international matches (2022+) ---
    # Filter to matches involving at least one 2026 WC team so the parameter
    # space stays tractable (prevents optimiser from managing 270+ teams).
    wc_teams = _get_wc2026_team_names()
    recent = fetch_recent_international_data()
    if not recent.empty and wc_teams:
        in_wc = recent["home_team"].isin(wc_teams) | recent["away_team"].isin(wc_teams)
        recent = recent[in_wc]

    if not recent.empty:
        rows_recent = pd.DataFrame({
            "date": recent["date"],
            "home_team": recent["home_team"],
            "away_team": recent["away_team"],
            "home_score": recent["home_score"],
            "away_score": recent["away_score"],
            "base_weight": recent["tournament"].apply(_recent_match_weight),
        })
        all_rows = pd.concat([rows_hist, rows_recent], ignore_index=True)
        n_recent = len(rows_recent)
    else:
        all_rows = rows_hist
        n_recent = 0

    # --- Temporal decay (applied to full combined dataset) ---
    ref_date = all_rows["date"].max()
    all_rows["weeks_ago"] = (ref_date - all_rows["date"]).dt.days / 7.0
    all_rows["weight"] = all_rows["base_weight"] * np.exp(-_DECAY_XI * all_rows["weeks_ago"])

    # --- Team index (union of all teams seen) ---
    all_teams = sorted(
        set(all_rows["home_team"].tolist()) | set(all_rows["away_team"].tolist())
    )
    team_idx: dict[str, int] = {t: i for i, t in enumerate(all_teams)}

    # Drop rows where either team name failed to map (rare edge case)
    mask = all_rows["home_team"].isin(team_idx) & all_rows["away_team"].isin(team_idx)
    all_rows = all_rows[mask]

    home_t = np.array([team_idx[t] for t in all_rows["home_team"]])
    away_t = np.array([team_idx[t] for t in all_rows["away_team"]])
    home_g = all_rows["home_score"].to_numpy()
    away_g = all_rows["away_score"].to_numpy()
    weights = all_rows["weight"].to_numpy()

    return {
        "home_t": home_t,
        "away_t": away_t,
        "home_g": home_g,
        "away_g": away_g,
        "weights": weights,
        "team_idx": team_idx,
        "all_teams": all_teams,
        "n_teams": len(all_teams),
        "n_matches_wc": len(rows_hist),
        "n_matches_recent": n_recent,
    }


# ── model fitting ─────────────────────────────────────────────────────────────

def fit_model(data: dict, prior_means: dict[str, float]) -> Any:
    """Fit Dixon-Coles Poisson model via MAP with ranking-informed Gaussian priors.

    prior_means: dict mapping team name → prior mean for total strength (α+δ).
    Each is split evenly: atk_prior[i] = dfs_prior[i] = prior_means[team] / 2.
    """
    n = data["n_teams"]
    all_teams = data["all_teams"]
    home_t = data["home_t"]
    away_t = data["away_t"]
    home_g = data["home_g"]
    away_g = data["away_g"]
    weights = data["weights"]

    # Build prior mean arrays (split total strength 50/50 between atk and dfs)
    strength_prior = np.array([prior_means.get(t, 0.0) for t in all_teams])
    atk_prior = strength_prior / 2.0
    dfs_prior = strength_prior / 2.0

    def nll(params: np.ndarray) -> float:
        mu = params[0]
        gamma = params[1]
        atk = params[2:2 + n]
        dfs = params[2 + n:]

        lh = np.exp(np.clip(mu + atk[home_t] - dfs[away_t] + gamma, -5, 5))
        la = np.exp(np.clip(mu + atk[away_t] - dfs[home_t], -5, 5))

        ll = weights * (
            home_g * np.log(np.maximum(lh, 1e-9)) - lh
            + away_g * np.log(np.maximum(la, 1e-9)) - la
        )
        # Ridge around FIFA-ranking-informed prior means (not around zero)
        ridge = _RIDGE_LAMBDA * 0.5 * (
            np.sum((atk - atk_prior) ** 2) + np.sum((dfs - dfs_prior) ** 2)
        )
        return -ll.sum() + ridge

    # Initialise at prior means rather than zero
    x0 = np.zeros(2 + 2 * n)
    x0[0] = math.log(1.5)
    x0[1] = 0.3
    x0[2:2 + n] = atk_prior
    x0[2 + n:] = dfs_prior

    return minimize(
        nll,
        x0,
        method="L-BFGS-B",
        options={"maxiter": 5000, "ftol": 1e-10, "gtol": 1e-7, "maxfun": 100000},
    )


def extract_parameters(result: Any, n_teams: int) -> dict:
    """Unpack optimiser result into named arrays with Laplace uncertainties."""
    params = result.x
    try:
        hess_inv_diag = np.diag(np.array(result.hess_inv.todense()))
    except Exception:
        hess_inv_diag = np.ones(len(params)) * 0.01

    std = np.sqrt(np.abs(hess_inv_diag))

    return {
        "mu": float(params[0]),
        "gamma": float(params[1]),
        "attack": params[2:2 + n_teams],
        "defense": params[2 + n_teams:],
        "attack_std": std[2:2 + n_teams],
        "defense_std": std[2 + n_teams:],
        "converged": bool(result.success),
        "n_iter": int(result.nit),
    }


# ── match prediction ──────────────────────────────────────────────────────────

def _dc_correction(hg: int, ag: int, lh: float, la: float, rho: float) -> float:
    if hg == 0 and ag == 0:
        return 1.0 - lh * la * rho
    if hg == 1 and ag == 0:
        return 1.0 + la * rho
    if hg == 0 and ag == 1:
        return 1.0 + lh * rho
    if hg == 1 and ag == 1:
        return 1.0 - rho
    return 1.0


def predict_match_probs(
    team_a: str,
    team_b: str,
    team_idx: dict[str, int],
    params: dict,
) -> dict:
    """Predict win/draw/loss probabilities for team_a (home) vs team_b (away)."""
    a_known = team_a in team_idx
    b_known = team_b in team_idx

    atk_a = params["attack"][team_idx[team_a]] if a_known else 0.0
    dfs_a = params["defense"][team_idx[team_a]] if a_known else 0.0
    atk_b = params["attack"][team_idx[team_b]] if b_known else 0.0
    dfs_b = params["defense"][team_idx[team_b]] if b_known else 0.0

    mu = params["mu"]
    gamma = params["gamma"]

    lh = math.exp(mu + atk_a - dfs_b + gamma)
    la = math.exp(mu + atk_b - dfs_a)

    G = _MAX_GOALS + 1
    pmf_h = np.array([poisson.pmf(g, lh) for g in range(G)])
    pmf_a = np.array([poisson.pmf(g, la) for g in range(G)])
    mat = np.outer(pmf_h, pmf_a)

    for hg in range(2):
        for ag in range(2):
            mat[hg, ag] *= max(_dc_correction(hg, ag, lh, la, _DC_RHO), 0.0)

    mat = mat / mat.sum()

    home_prob = float(np.tril(mat, -1).sum())
    draw_prob = float(np.trace(mat))
    away_prob = float(np.triu(mat, 1).sum())

    total = home_prob + draw_prob + away_prob
    home_prob /= total
    draw_prob /= total
    away_prob /= total

    if home_prob >= draw_prob and home_prob >= away_prob:
        predicted = "HOME"
    elif away_prob >= draw_prob and away_prob >= home_prob:
        predicted = "AWAY"
    else:
        predicted = "DRAW"

    return {
        "home_prob": home_prob,
        "draw_prob": draw_prob,
        "away_prob": away_prob,
        "predicted": predicted,
        "lambda_home": lh,
        "lambda_away": la,
        "team_a_known": a_known,
        "team_b_known": b_known,
    }


def predict_all_2026_matches(params: dict, team_idx: dict[str, int]) -> list[dict]:
    matches = get_all_named_matches()
    results = []
    for m in matches:
        hist_a = _to_hist_name(m["team1"])
        hist_b = _to_hist_name(m["team2"])
        probs = predict_match_probs(hist_a, hist_b, team_idx, params)
        results.append({
            "team1": m["team1"],
            "team2": m["team2"],
            "hist_team1": hist_a,
            "hist_team2": hist_b,
            "date": m.get("date", ""),
            "group": m.get("group", m.get("round", "")),
            "score": m.get("score"),
            **probs,
        })
    return results


def compute_bayesian_calibration(params: dict, team_idx: dict[str, int]) -> dict:
    from benchmark.runner import _load_results

    raw = _load_results()
    if not raw:
        return {"brier": None, "log_loss": None, "accuracy": None, "n_matches": 0}

    seen: set[str] = set()
    played: list[dict] = []
    for r in raw:
        key = f"{r['team1']}||{r['team2']}||{r['date']}"
        if key not in seen and r.get("actual_result") not in (None, "", "UNKNOWN"):
            seen.add(key)
            played.append(r)

    if not played:
        return {"brier": None, "log_loss": None, "accuracy": None, "n_matches": 0}

    brier_sum = 0.0
    log_loss_sum = 0.0
    correct = 0

    for r in played:
        hist_a = _to_hist_name(r["team1"])
        hist_b = _to_hist_name(r["team2"])
        probs = predict_match_probs(hist_a, hist_b, team_idx, params)

        actual = r["actual_result"]
        p_actual = probs.get(f"{actual.lower()}_prob", 0.0)
        brier_sum += (p_actual - 1.0) ** 2
        log_loss_sum += -math.log(max(p_actual, 1e-6))
        if probs["predicted"] == actual:
            correct += 1

    n = len(played)
    return {
        "brier": brier_sum / n,
        "log_loss": log_loss_sum / n,
        "accuracy": correct / n,
        "n_matches": n,
    }


# ── top-level entry point ─────────────────────────────────────────────────────

def fit_and_predict() -> dict:
    """Fit the model and predict all 2026 matches. Result is cached after first call."""
    global _model_cache
    if _model_cache is not None:
        return _model_cache

    data = load_training_data()
    prior_means = get_prior_means(data["all_teams"])
    result = fit_model(data, prior_means)
    params = extract_parameters(result, data["n_teams"])
    predictions = predict_all_2026_matches(params, data["team_idx"])
    calibration = compute_bayesian_calibration(params, data["team_idx"])

    _model_cache = {
        "params": params,
        "team_idx": data["team_idx"],
        "all_teams": data["all_teams"],
        "n_matches_wc": data["n_matches_wc"],
        "n_matches_recent": data["n_matches_recent"],
        "predictions": predictions,
        "calibration": calibration,
        "converged": params["converged"],
    }
    return _model_cache