utils/scoring.py

"""Scoring rules for the WC 2026 prediction game."""
from __future__ import annotations

RULES: dict[str, int] = {
    # Group stage
    "group_match_correct": 1,       # correct W/D/L outcome
    "group_winner": 3,              # correct group winner (1st place)
    "group_runner_up": 2,           # correct group 2nd place
    # Knockout — points for correctly predicting the advancing team
    "r32_correct": 2,
    "r16_correct": 3,
    "qf_correct": 4,
    "sf_correct": 5,
    "finalist_correct": 8,          # reaching the final
    "champion_correct": 15,
    # Bonus
    "both_finalists_correct": 5,    # both finalists right (order-independent)
    "exact_final_result": 3,        # correct final score outcome (W for which team)
}

# Stage name → scoring key
_STAGE_KEY: dict[str, str] = {
    "ROUND_OF_32": "r32_correct",
    "ROUND_OF_16": "r16_correct",
    "QUARTER_FINALS": "qf_correct",
    "SEMI_FINALS": "sf_correct",
    "THIRD_PLACE": "sf_correct",    # same value as sf
    "FINAL": "finalist_correct",
}


def compute_score(predictions: list[dict], results: list[dict]) -> dict:
    """Compute score for a single model/player's predictions.

    Args:
        predictions: list of prediction dicts with keys game_id, predicted_result,
                     home_team, away_team, stage, group_id.
        results: list of actual result dicts with keys game_id, actual_result
                 (same encoding: 'A Win', 'Draw', 'B Win'), home_team, away_team,
                 winner (team name of advancing team for KO rounds).

    Returns:
        dict with keys 'total', 'breakdown' (dict of category -> points).
    """
    if not results:
        return {"total": 0, "breakdown": {k: 0 for k in RULES}}

    results_by_id: dict[str, dict] = {r["game_id"]: r for r in results}
    pred_by_id: dict[str, dict] = {p["game_id"]: p for p in predictions}

    breakdown: dict[str, int] = {k: 0 for k in RULES}

    for game_id, pred in pred_by_id.items():
        actual = results_by_id.get(game_id)
        if actual is None:
            continue

        stage = str(pred.get("stage", "")).upper()

        if stage == "GROUP_STAGE":
            if pred.get("predicted_result") == actual.get("actual_result"):
                breakdown["group_match_correct"] += RULES["group_match_correct"]
        else:
            # Knockout: check if predicted advancing team matches actual winner
            key = _STAGE_KEY.get(stage)
            if key:
                pred_winner = _resolve_ko_winner(pred)
                actual_winner = actual.get("winner", "")
                if pred_winner and pred_winner == actual_winner:
                    breakdown[key] += RULES[key]

    # Group standings bonuses
    group_results = [r for r in results if str(r.get("stage", "")).upper() == "GROUP_STAGE"]
    if group_results:
        _score_group_standings(pred_by_id, results_by_id, group_results, breakdown)

    # Final bonuses
    final_results = [r for r in results if str(r.get("stage", "")).upper() == "FINAL"]
    if final_results and final_results[0].get("game_id") in pred_by_id:
        final_res = final_results[0]
        final_pred = pred_by_id[final_res["game_id"]]
        # Champion bonus
        actual_champ = final_res.get("winner", "")
        pred_champ = _resolve_ko_winner(final_pred)
        if pred_champ == actual_champ:
            breakdown["champion_correct"] += RULES["champion_correct"]
        # Both finalists
        actual_finalists = {final_res.get("home_team", ""), final_res.get("away_team", "")}
        pred_finalists = {final_pred.get("home_team", ""), final_pred.get("away_team", "")}
        if actual_finalists == pred_finalists:
            breakdown["both_finalists_correct"] += RULES["both_finalists_correct"]
        # Exact final result (which team wins)
        if final_pred.get("predicted_result") == final_res.get("actual_result"):
            breakdown["exact_final_result"] += RULES["exact_final_result"]

    total = sum(breakdown.values())
    return {"total": total, "breakdown": breakdown}


def _resolve_ko_winner(pred: dict) -> str:
    """Determine the predicted advancing team for a knockout match."""
    result = pred.get("predicted_result", "")
    if result == "A Win":
        return str(pred.get("home_team", ""))
    elif result == "B Win":
        return str(pred.get("away_team", ""))
    # Draw in KO → no clear winner prediction (could be penalties)
    return ""


def _score_group_standings(
    pred_by_id: dict[str, dict],
    results_by_id: dict[str, dict],
    group_results: list[dict],
    breakdown: dict[str, int],
) -> None:
    """Award group winner/runner-up bonus points.

    Only scores a group once *all* of its matches have been played, so partial
    standings after 1-2 games never award a (premature) winner/runner-up bonus.
    """
    # Expected matches per group (from the full prediction set) vs. played so far.
    expected_per_group: dict[str, int] = {}
    for pred in pred_by_id.values():
        if str(pred.get("stage", "")).upper() != "GROUP_STAGE":
            continue
        gid = pred.get("group_id", "")
        if gid:
            expected_per_group[gid] = expected_per_group.get(gid, 0) + 1

    played_per_group: dict[str, int] = {}
    for r in group_results:
        gid = r.get("group_id", "")
        if gid:
            played_per_group[gid] = played_per_group.get(gid, 0) + 1

    complete_groups = {
        gid
        for gid, expected in expected_per_group.items()
        if expected > 0 and played_per_group.get(gid, 0) >= expected
    }

    # Collect actual group standings
    actual_groups: dict[str, dict[str, int]] = {}
    for r in group_results:
        gid = r.get("group_id", "")
        if not gid:
            continue
        for team_key in ("home_team", "away_team"):
            team = r.get(team_key, "")
            if team:
                actual_groups.setdefault(gid, {})[team] = actual_groups.get(gid, {}).get(team, 0)
        actual_result = r.get("actual_result", "")
        home = r.get("home_team", "")
        away = r.get("away_team", "")
        if actual_result == "A Win":
            actual_groups[gid][home] = actual_groups[gid].get(home, 0) + 3
        elif actual_result == "Draw":
            actual_groups[gid][home] = actual_groups[gid].get(home, 0) + 1
            actual_groups[gid][away] = actual_groups[gid].get(away, 0) + 1
        elif actual_result == "B Win":
            actual_groups[gid][away] = actual_groups[gid].get(away, 0) + 3

    # Collect predicted group standings from predictions
    pred_groups: dict[str, dict[str, int]] = {}
    for game_id, pred in pred_by_id.items():
        if str(pred.get("stage", "")).upper() != "GROUP_STAGE":
            continue
        gid = pred.get("group_id", "")
        if not gid:
            continue
        pred_result = pred.get("predicted_result", "")
        home = pred.get("home_team", "")
        away = pred.get("away_team", "")
        if home:
            pred_groups.setdefault(gid, {})[home] = pred_groups.get(gid, {}).get(home, 0)
        if away:
            pred_groups.setdefault(gid, {})[away] = pred_groups.get(gid, {}).get(away, 0)
        if pred_result == "A Win":
            pred_groups[gid][home] = pred_groups[gid].get(home, 0) + 3
        elif pred_result == "Draw":
            pred_groups[gid][home] = pred_groups[gid].get(home, 0) + 1
            pred_groups[gid][away] = pred_groups[gid].get(away, 0) + 1
        elif pred_result == "B Win":
            pred_groups[gid][away] = pred_groups[gid].get(away, 0) + 3

    for gid in actual_groups:
        if gid not in pred_groups or gid not in complete_groups:
            continue
        actual_sorted = sorted(actual_groups[gid], key=lambda t: -actual_groups[gid][t])
        pred_sorted = sorted(pred_groups[gid], key=lambda t: -pred_groups[gid].get(t, 0))
        if len(actual_sorted) >= 1 and len(pred_sorted) >= 1:
            if actual_sorted[0] == pred_sorted[0]:
                breakdown["group_winner"] += RULES["group_winner"]
        if len(actual_sorted) >= 2 and len(pred_sorted) >= 2:
            if actual_sorted[1] == pred_sorted[1]:
                breakdown["group_runner_up"] += RULES["group_runner_up"]


def compute_all_scores(
    all_predictions: dict[str, list[dict]],
    results: list[dict],
) -> dict[str, dict]:
    """Compute scores for all models/players.

    Args:
        all_predictions: mapping of model_name → list of prediction dicts
        results: list of actual result dicts

    Returns:
        mapping of model_name → score dict (total + breakdown)
    """
    return {
        model_name: compute_score(preds, results)
        for model_name, preds in all_predictions.items()
    }


def compute_max_possible(predictions_template: list[dict], results: list[dict]) -> dict:
    """Maximum score attainable given the matches played so far.

    Builds an *oracle* from the full fixture template (so group-completeness gating
    sees the true match count per group), overriding the predicted outcome of every
    played match to the actual result / advancing team. The total is the ceiling any
    model could have reached, counting group winner/runner-up bonuses only once a
    group is complete.
    """
    if not results:
        return {"total": 0, "breakdown": {k: 0 for k in RULES}}

    def _oracle_result(r: dict) -> str:
        actual_result = r.get("actual_result", "")
        if actual_result:
            return actual_result
        winner = r.get("winner", "")
        if winner == r.get("home_team", ""):
            return "A Win"
        if winner == r.get("away_team", ""):
            return "B Win"
        return ""

    override = {r.get("game_id"): _oracle_result(r) for r in results}

    oracle: list[dict] = []
    for fixture in predictions_template:
        gid = fixture.get("game_id")
        oracle.append({**fixture, "predicted_result": override.get(gid, fixture.get("predicted_result", ""))})

    return compute_score(oracle, results)