arena.py

"""Match logic, Elo and the public results store (a HF dataset repo)."""
from __future__ import annotations

import io
import json
import os
import random
import threading
import uuid
from datetime import datetime, timezone

from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download

from tetris import TetrisGame

MAX_PIECES = int(os.environ.get("MAX_PIECES", 500))
K_FACTOR = 32
START_ELO = 1000.0
APP_VERSION = "1.2"
# Bump to start a fresh leaderboard: data saved under another season is ignored.
SEASON = 4


def choose(game: TetrisGame, player, protocol: str, seed: int):
    """Let `player` pick the next placement. Returns False on top-out."""
    cands = game.candidates()
    if not cands:
        game.top_out()
        return False
    vals = player.values(protocol, cands)
    best = max(vals)
    ties = [i for i, v in enumerate(vals) if v >= best - 1e-9]
    # deterministic tie-break, identical for every player at the same piece index
    rng = random.Random(f"{seed}:{game.pieces}")
    i = ties[rng.randrange(len(ties))]
    game.apply(cands[i], vals[i], len(cands))
    return True


def result_key(g: TetrisGame):
    return (g.score, g.lines, g.pieces)


def rank_games(games):
    """Competition ranking (1,1,3...) by score, then lines, then pieces survived."""
    keys = [result_key(g) for g in games]
    return [1 + sum(1 for k2 in keys if k2 > k) for k in keys]


def elo_deltas(ratings: list[float], keys: list[tuple]) -> list[float]:
    """Multiplayer Elo: every pair is a mini-game, scaled by 1/(N-1)."""
    n = len(ratings)
    out = []
    for i in range(n):
        s = 0.0
        for j in range(n):
            if i == j:
                continue
            actual = 1.0 if keys[i] > keys[j] else 0.5 if keys[i] == keys[j] else 0.0
            expected = 1.0 / (1.0 + 10 ** ((ratings[j] - ratings[i]) / 400.0))
            s += actual - expected
        out.append(K_FACTOR * s / (n - 1))
    return out


class ResultsStore:
    """Leaderboard persisted in a public dataset repo: leaderboard.json + one JSON per match."""

    def __init__(self, repo_id: str, token: str | None, min_params: int = 0):
        self.repo_id = repo_id
        self.token = token
        self.min_params = min_params
        self.api = HfApi(token=token)
        self.lock = threading.Lock()
        self.load_error = None
        self.save_error = None
        self.data = self._load()

    @property
    def persistent(self):
        return bool(self.token)

    def _empty(self):
        return {"version": 1, "season": SEASON, "updated": None, "protocols": {"guided": {}, "blind": {}}}

    def _load(self):
        try:
            path = hf_hub_download(self.repo_id, "leaderboard.json", repo_type="dataset", token=self.token, force_download=True)
            data = json.load(open(path))
            if data.get("season") != SEASON:
                self.load_error = None
                return self._empty()  # previous season: start clean
            data.setdefault("protocols", {}).setdefault("guided", {})
            data["protocols"].setdefault("blind", {})
            self._prune(data)
            self.load_error = None
            return data
        except Exception as e:  # first run (no file yet) or Hub unreachable
            self.load_error = f"leaderboard not loaded ({type(e).__name__}); starting empty"
            return self._empty()

    def _prune(self, data):
        """Drop models below the parameter minimum (rule added after they were rated)."""
        for table in data["protocols"].values():
            for model_id in [m for m, e in table.items() if 0 < (e.get("params") or 0) < self.min_params]:
                del table[model_id]

    def reload(self):
        with self.lock:
            self.data = self._load()

    def rating(self, protocol, model_id):
        e = self.data["protocols"][protocol].get(model_id)
        return e["elo"] if e else START_ELO

    def record(self, protocol: str, seed: int, players, games) -> dict:
        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
        with self.lock:
            table = self.data["protocols"][protocol]
            ratings = [self.rating(protocol, p.model_id) for p in players]
            keys = [result_key(g) for g in games]
            deltas = elo_deltas(ratings, keys)
            ranks = rank_games(games)
            match = {
                "id": uuid.uuid4().hex[:12],
                "time": now,
                "protocol": protocol,
                "seed": seed,
                "max_pieces": MAX_PIECES,
                "app_version": APP_VERSION,
                "season": SEASON,
                "players": [],
            }
            for p, g, r, d, rank in zip(players, games, ratings, deltas, ranks):
                e = table.setdefault(
                    p.model_id,
                    {"model": p.model_id, "elo": START_ELO, "games": 0, "wins": 0, "total_pieces": 0,
                     "total_lines": 0, "best_score": 0, "best_lines": 0},
                )
                e["elo"] = round(r + d, 2)
                e["games"] += 1
                e["wins"] += int(rank == 1)
                e["total_pieces"] += g.pieces
                e["total_lines"] += g.lines
                e["best_score"] = max(e["best_score"], g.score)
                e["best_lines"] = max(e["best_lines"], g.lines)
                e["params"] = p.n_params
                e["custom_code"] = p.custom_code
                e["last_sha"] = p.sha
                e["last_played"] = now
                match["players"].append({
                    "model": p.model_id, "sha": p.sha, "params": p.n_params, "rank": rank,
                    "score": g.score, "lines": g.lines, "pieces": g.pieces, "tetrises": g.tetrises,
                    "survived": g.alive, "elo_before": round(r, 2), "elo_after": e["elo"], "elo_delta": round(d, 2),
                })
            self.data["updated"] = now
            if self.persistent:
                try:
                    day = now[:10]
                    self.api.create_commit(
                        repo_id=self.repo_id,
                        repo_type="dataset",
                        operations=[
                            CommitOperationAdd("leaderboard.json", io.BytesIO(json.dumps(self.data, indent=1).encode())),
                            CommitOperationAdd(f"matches/{day}/{match['id']}.json", io.BytesIO(json.dumps(match, indent=1).encode())),
                        ],
                        commit_message=f"{protocol} match {match['id']}: " + " vs ".join(p.model_id for p in players),
                    )
                    self.save_error = None
                except Exception as e:
                    self.save_error = f"could not save to {self.repo_id} ({type(e).__name__}: {str(e)[:120]})"
            return match

    def rows(self, protocol: str):
        entries = sorted(self.data["protocols"][protocol].values(), key=lambda e: -e["elo"])
        return entries