drama.py

from __future__ import annotations

from collections import defaultdict

import numpy as np
from sklearn.ensemble import HistGradientBoostingRegressor
from sklearn.inspection import permutation_importance
from sklearn.metrics import mean_absolute_error, r2_score


FEATURES = (
    "stage pressure",
    "combined points per game",
    "form gap",
    "goals scored per game",
    "goals allowed per game",
    "fouls per game",
    "cards per game",
    "shots per game",
    "tournament experience",
)

STAGES = {
    "group-stage": 0.15,
    "round-of-32": 0.3,
    "round-of-16": 0.48,
    "quarterfinals": 0.68,
    "semifinals": 0.86,
    "3rd-place-match": 0.5,
    "final": 1.0,
}

STAGE_NAMES = {
    "group-stage": "Group stage",
    "round-of-32": "Round of 32",
    "round-of-16": "Round of 16",
    "quarterfinals": "Quarter-final",
    "semifinals": "Semi-final",
    "3rd-place-match": "Third-place match",
    "final": "Final",
}


def clamp(value: float) -> int:
    return max(0, min(100, round(value)))


def competitors(event: dict) -> tuple[dict, dict]:
    teams = event["competitions"][0]["competitors"]
    ordered = sorted(teams, key=lambda item: item["homeAway"] != "home")
    return ordered[0], ordered[1]


def stat(team: dict, name: str) -> float:
    for item in team.get("statistics", []):
        if item.get("name") == name:
            try:
                return float(str(item.get("displayValue", 0)).replace("%", ""))
            except ValueError:
                return 0.0
    return 0.0


def card_counts(event: dict) -> tuple[int, int]:
    yellow = red = 0
    for detail in event["competitions"][0].get("details", []):
        kind = detail.get("type", {}).get("text", "").lower()
        yellow += "yellow card" in kind
        red += "red card" in kind
    return yellow, red


def team_cards(event: dict, team_id: str) -> int:
    total = 0
    for detail in event["competitions"][0].get("details", []):
        if str(detail.get("team", {}).get("id")) != str(team_id):
            continue
        kind = detail.get("type", {}).get("text", "").lower()
        total += 1 if "yellow card" in kind else 2 if "red card" in kind else 0
    return total


def observed_index(event: dict) -> int:
    home, away = competitors(event)
    yellow, red = card_counts(event)
    fouls = stat(home, "foulsCommitted") + stat(away, "foulsCommitted")
    goals = int(float(home.get("score", 0))) + int(float(away.get("score", 0)))
    score_gap = abs(int(float(home.get("score", 0))) - int(float(away.get("score", 0))))
    details = event["competitions"][0].get("details", [])
    late_goal = any(
        "goal" in detail.get("type", {}).get("text", "").lower()
        and float(detail.get("clock", {}).get("value", 0)) >= 75 * 60
        for detail in details
    )
    status = event["status"]["type"].get("description", "").lower()
    knockout = STAGES.get(event.get("season", {}).get("slug", ""), 0.15)
    raw = 10 + fouls * 0.78 + yellow * 4.8 + red * 12 + min(goals, 6) * 3.2
    raw += (10 if score_gap <= 1 else 0) + (9 if late_goal else 0) + knockout * 12
    raw += 10 if "extra time" in status or "penalties" in status else 0
    return clamp(raw)


def _empty_form() -> dict:
    return {"games": 0, "points": 0, "gf": 0, "ga": 0, "fouls": 0.0, "cards": 0, "shots": 0.0}


def _rate(form: dict, key: str, fallback: float) -> float:
    return form[key] / form["games"] if form["games"] else fallback


def _features(event: dict, forms: dict[str, dict]) -> tuple[list[float], dict]:
    home, away = competitors(event)
    a = forms[home["team"]["id"]]
    b = forms[away["team"]["id"]]
    ppg_a, ppg_b = _rate(a, "points", 1.5), _rate(b, "points", 1.5)
    values = [
        STAGES.get(event.get("season", {}).get("slug", ""), 0.15),
        (ppg_a + ppg_b) / 6,
        abs(ppg_a - ppg_b) / 3,
        (_rate(a, "gf", 1.25) + _rate(b, "gf", 1.25)) / 6,
        (_rate(a, "ga", 1.25) + _rate(b, "ga", 1.25)) / 6,
        (_rate(a, "fouls", 11.5) + _rate(b, "fouls", 11.5)) / 35,
        (_rate(a, "cards", 1.8) + _rate(b, "cards", 1.8)) / 8,
        (_rate(a, "shots", 10) + _rate(b, "shots", 10)) / 35,
        min((a["games"] + b["games"]) / 10, 1),
    ]
    raw = {
        "stage": STAGE_NAMES.get(event.get("season", {}).get("slug", ""), "Tournament match"),
        "home_ppg": round(ppg_a, 2),
        "away_ppg": round(ppg_b, 2),
        "combined_fouls_pg": round(_rate(a, "fouls", 11.5) + _rate(b, "fouls", 11.5), 1),
        "combined_cards_pg": round(_rate(a, "cards", 1.8) + _rate(b, "cards", 1.8), 1),
        "combined_shots_pg": round(_rate(a, "shots", 10) + _rate(b, "shots", 10), 1),
        "prior_games": a["games"] + b["games"],
    }
    return values, raw


def _update(forms: dict[str, dict], event: dict) -> None:
    home, away = competitors(event)
    home_score, away_score = int(float(home["score"])), int(float(away["score"]))
    for team, scored, allowed in (
        (home, home_score, away_score),
        (away, away_score, home_score),
    ):
        form = forms[team["team"]["id"]]
        form["games"] += 1
        form["points"] += 3 if scored > allowed else 1 if scored == allowed else 0
        form["gf"] += scored
        form["ga"] += allowed
        form["fouls"] += stat(team, "foulsCommitted")
        form["cards"] += team_cards(event, team["team"]["id"])
        form["shots"] += stat(team, "totalShots")


def train_and_score(events: list[dict], target: dict) -> dict:
    ordered = sorted(events, key=lambda item: item["date"])
    forms: dict[str, dict] = defaultdict(_empty_form)
    x: list[list[float]] = []
    y: list[int] = []
    target_x = target_raw = None

    for event in ordered:
        features, raw = _features(event, forms)
        if str(event["id"]) == str(target["id"]):
            target_x, target_raw = features, raw
            continue
        if event["date"] > target["date"]:
            continue
        if event["status"]["type"].get("completed"):
            x.append(features)
            y.append(observed_index(event))
            _update(forms, event)

    if target_x is None:
        target_x, target_raw = _features(target, forms)

    x_data, y_data = np.asarray(x), np.asarray(y)
    split = max(30, int(len(x_data) * 0.8))
    eval_model = HistGradientBoostingRegressor(
        max_iter=120, max_leaf_nodes=8, l2_regularization=2, random_state=26
    )
    eval_model.fit(x_data[:split], y_data[:split])
    held_out = eval_model.predict(x_data[split:])
    mae = mean_absolute_error(y_data[split:], held_out)
    r2 = r2_score(y_data[split:], held_out)
    baseline_mae = mean_absolute_error(y_data[split:], np.full(len(y_data[split:]), y_data[:split].mean()))
    baseline_lift = (baseline_mae - mae) / baseline_mae * 100

    model = HistGradientBoostingRegressor(
        max_iter=160, max_leaf_nodes=8, l2_regularization=2, random_state=26
    )
    model.fit(x_data, y_data)
    forecast = clamp(model.predict(np.asarray([target_x]))[0])
    permuted = permutation_importance(model, x_data, y_data, n_repeats=8, random_state=26)
    weights = np.maximum(permuted.importances_mean, 0)
    weights = weights / weights.sum() if weights.sum() else np.ones(len(FEATURES)) / len(FEATURES)
    importances = sorted(zip(FEATURES, weights), key=lambda item: item[1], reverse=True)
    confidence = clamp(88 - mae * 1.6 + min(target_raw["prior_games"], 10))
    return {
        "forecast": forecast,
        "confidence": confidence,
        "samples": len(x),
        "mae": round(float(mae), 1),
        "r2": round(float(r2), 2),
        "baseline_lift": round(float(baseline_lift), 1),
        "features": target_raw,
        "importances": [(name, round(float(weight) * 100, 1)) for name, weight in importances],
    }


def h2h(summary: dict) -> dict:
    groups = summary.get("headToHeadGames", [])
    games = groups[0].get("events", []) if groups else []
    world_cups = [game for game in games if "World Cup" in game.get("leagueName", "")]
    shootouts = [game for game in games if int(game.get("homeShootoutScore", 0)) or int(game.get("awayShootoutScore", 0))]
    latest = games[0] if games else None
    return {
        "games": len(games),
        "world_cups": len(world_cups),
        "shootouts": len(shootouts),
        "latest": f"{latest['competitionName']} · {latest['score']}" if latest else "No H2H record returned",
    }