Teaching a Neural Network to Play Go

import math, random, time
from collections import defaultdict

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

SIZE = 5          # board edge — the one number that changes the whole project
KOMI = 0.0        # White's compensation; 0 is the setting the solved 5x5 result uses
N_POINTS = SIZE * SIZE
N_MOVES = N_POINTS + 1     # every intersection, plus pass

EMPTY, BLACK, WHITE, PASS = 0, 1, 2, -1
OTHER = {BLACK: WHITE, WHITE: BLACK}

DEVICE = "mps" if torch.backends.mps.is_available() else "cpu"
torch.manual_seed(0); random.seed(0); np.random.seed(0)

# Where the expensive artefacts live. On a fresh clone these files do not exist
# and every cell below computes them; in this repository they are already here,
# so the lab loads exactly the datasets and checkpoints the note's numbers came
# from instead of a hastily shrunk imitation of them.
from pathlib import Path
ART = Path("scripts/fixtures")
CKPT = Path("scripts/checkpoints")

def cached(path, build, label=""):
    """Load `path` if it exists, otherwise build it and save it there."""
    path = Path(path)
    if path.exists():
        print(f"loading {path}{' — ' + label if label else ''}")
        return np.load(path)
    print(f"computing {path}{' — ' + label if label else ''} (this is the slow part)")
    data = build()
    path.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(path, **{k: v for k, v in data.items()
                                 if isinstance(v, np.ndarray)})
    return np.load(path)

print("device:", DEVICE)

def neighbours(size):
    orth, diag = [], []
    for p in range(size * size):
        r, c = divmod(p, size)
        orth.append(tuple(rr * size + cc for rr, cc in
                          ((r-1, c), (r+1, c), (r, c-1), (r, c+1))
                          if 0 <= rr < size and 0 <= cc < size))
        diag.append(tuple(rr * size + cc for rr, cc in
                          ((r-1, c-1), (r-1, c+1), (r+1, c-1), (r+1, c+1))
                          if 0 <= rr < size and 0 <= cc < size))
    return orth, diag

ORTH, DIAG = neighbours(SIZE)


class Go:
    """A position: stones, side to move, and every board seen so far (for superko)."""

    __slots__ = ("points", "to_play", "history", "n_passes", "move_count", "captures")

    def __init__(self):
        self.points = bytearray(N_POINTS)
        self.to_play = BLACK
        self.history = frozenset({bytes(self.points)})
        self.n_passes = 0
        self.move_count = 0
        self.captures = {BLACK: 0, WHITE: 0}

    def copy(self):
        new = Go.__new__(Go)
        new.points = bytearray(self.points)
        new.to_play = self.to_play
        new.history = self.history
        new.n_passes = self.n_passes
        new.move_count = self.move_count
        new.captures = dict(self.captures)
        return new

    # --- groups -----------------------------------------------------------
    def _has_liberty(self, points, p):
        colour = points[p]
        stack, seen = [p], {p}
        while stack:
            q = stack.pop()
            for n in ORTH[q]:
                if points[n] == EMPTY:
                    return True
                if points[n] == colour and n not in seen:
                    seen.add(n); stack.append(n)
        return False

    def _remove(self, points, p):
        colour = points[p]
        stack, seen, removed = [p], {p}, 0
        while stack:
            q = stack.pop()
            points[q] = EMPTY; removed += 1
            for n in ORTH[q]:
                if points[n] == colour and n not in seen:
                    seen.add(n); stack.append(n)
        return removed

    def group_liberties(self, p):
        """(stones in the group, number of its liberties) — for inspection."""
        colour = self.points[p]
        if colour == EMPTY:
            return [], 0
        stack, seen, group, libs = [p], {p}, [], set()
        while stack:
            q = stack.pop(); group.append(q)
            for n in ORTH[q]:
                if self.points[n] == EMPTY:
                    libs.add(n)
                elif self.points[n] == colour and n not in seen:
                    seen.add(n); stack.append(n)
        return group, len(libs)

    def _try(self, p, colour):
        """Resulting board and capture count, or None if the move is suicide."""
        board = bytearray(self.points)
        board[p] = colour
        captured = 0
        for n in ORTH[p]:
            if board[n] == OTHER[colour] and not self._has_liberty(board, n):
                captured += self._remove(board, n)
        if captured == 0 and not self._has_liberty(board, p):
            return None
        return board, captured

    # --- legality and play ------------------------------------------------
    def is_legal(self, move):
        if move == PASS:
            return True
        if not (0 <= move < N_POINTS) or self.points[move] != EMPTY:
            return False
        out = self._try(move, self.to_play)
        return out is not None and bytes(out[0]) not in self.history

    def play(self, move):
        new = self.copy()
        new.move_count += 1
        new.to_play = OTHER[self.to_play]
        if move == PASS:
            new.n_passes = self.n_passes + 1
            return new
        board, captured = self._try(move, self.to_play)
        new.points = board
        new.captures[self.to_play] += captured
        new.history = self.history | {bytes(board)}
        new.n_passes = 0
        return new

    def eye_like(self, p, colour):
        """Heuristic eye test: what stops a player filling in its own eyes."""
        if self.points[p] != EMPTY:
            return False
        if any(self.points[n] != colour for n in ORTH[p]):
            return False
        d = DIAG[p]
        friendly = sum(1 for n in d if self.points[n] == colour)
        return friendly == len(d) if len(d) < 4 else friendly >= 3

    def agent_moves(self):
        """What a player should consider: never fill your own eye, pass only if stuck.

        Under area scoring a stone inside your own territory still counts as your
        area, so playing on is never worse than passing while a move exists.
        Passing early is a mistake, not an option — letting search spend visits on
        it just pollutes the result.
        """
        moves = [p for p in range(N_POINTS)
                 if self.points[p] == EMPTY and not self.eye_like(p, self.to_play)
                 and self.is_legal(p)]
        return moves if moves else [PASS]

    # --- ending and scoring -----------------------------------------------
    def is_over(self):
        return self.n_passes >= 2

    def area_score(self):
        black = sum(1 for v in self.points if v == BLACK)
        white = sum(1 for v in self.points if v == WHITE)
        seen = set()
        for start in range(N_POINTS):
            if self.points[start] != EMPTY or start in seen:
                continue
            region, border, stack = [], set(), [start]
            seen.add(start)
            while stack:
                q = stack.pop(); region.append(q)
                for n in ORTH[q]:
                    if self.points[n] == EMPTY:
                        if n not in seen:
                            seen.add(n); stack.append(n)
                    else:
                        border.add(self.points[n])
            if border == {BLACK}:
                black += len(region)
            elif border == {WHITE}:
                white += len(region)
        return black, white

    def result(self):
        b, w = self.area_score()
        return b - (w + KOMI)

    def winner(self):
        return BLACK if self.result() > 0 else WHITE

    def outcome_for(self, colour):
        return 1.0 if self.winner() == colour else -1.0

    def __str__(self):
        g = {EMPTY: ".", BLACK: "X", WHITE: "O"}
        rows = [" ".join(g[v] for v in self.points[r*SIZE:(r+1)*SIZE]) for r in range(SIZE)]
        return "\n".join(rows)

def place(black=(), white=(), to_play=BLACK):
    """Build a position directly — for tests and demonstrations."""
    g = Go()
    letters = "ABCDE"
    for label, colour in [(l, BLACK) for l in black] + [(l, WHITE) for l in white]:
        p = (SIZE - int(label[1:])) * SIZE + letters.index(label[0])
        g.points[p] = colour
    g.to_play = to_play
    g.history = frozenset({bytes(g.points)})
    return g

ALL = [f"{c}{r}" for r in range(1, 6) for c in "ABCDE"]

def idx(label):
    """'C3' -> flat index, so the checks below read as Go rather than as numbers."""
    return (SIZE - int(label[1:])) * SIZE + "ABCDE".index(label[0])

# One eye: White can play the single liberty and take all 24 stones.
one_eye = place(black=[l for l in ALL if l != "C3"], to_play=WHITE)
assert one_eye.is_legal(idx("C3"))
after = one_eye.play(idx("C3"))
assert after.captures[WHITE] == 24

# Two eyes: White has no legal move at all — the group cannot be captured.
two_eyes = place(black=[l for l in ALL if l not in ("A1", "E5")], to_play=WHITE)
assert two_eyes.agent_moves() == [PASS]

# Suicide is illegal; the same point becomes legal when the move captures first.
assert not place(white=["A4", "B5"], to_play=BLACK).is_legal(idx("A5"))
assert place(black=["C5", "B4", "A3"], white=["B5", "A4"], to_play=BLACK).is_legal(idx("A5"))

# A genuine ko: the recapture is banned because it repeats a position.
ko = place(black=["B5", "A4", "B3"], white=["C5", "B4", "D4", "C3"], to_play=BLACK)
ko_after = ko.play(idx("C4"))                 # takes the lone white stone at B4
assert ko_after.captures[BLACK] == 1          # exactly one stone: a ko, not a net
assert not ko_after.is_legal(idx("B4"))       # recapturing would repeat a position

print("rules pass\n")
print("one eye — White to play, 24 black stones about to die:")
print(one_eye, "\n")
print("after White plays the eye:")
print(after)

def rollout_move(state, rng):
    """Uniform over sensible moves, by rejection — cheaper than listing them all."""
    empties = [p for p in range(N_POINTS) if state.points[p] == EMPTY]
    rng.shuffle(empties)
    colour = state.to_play
    for p in empties:
        if not state.eye_like(p, colour) and state.is_legal(p):
            return p
    return PASS

def playout(state, rng, max_moves=400):
    while not state.is_over() and state.move_count < max_moves:
        state = state.play(rollout_move(state, rng))
    return state

rng = random.Random(0)
t0 = time.time()
finals = [playout(Go(), rng) for _ in range(300)]
lengths = [f.move_count for f in finals]
stones = [sum(1 for v in f.points if v) for f in finals]
margins = [abs(f.result()) for f in finals]

print(f"{len(finals)} random playouts in {time.time()-t0:.1f}s")
print(f"length      {np.mean(lengths):.1f} moves (min {min(lengths)})")
print(f"stones left {np.mean(stones):.1f} / {N_POINTS}")
print(f"black wins  {np.mean([f.winner()==BLACK for f in finals]):.2f}")
print(f"mean |margin| {np.mean(margins):.1f} of {N_POINTS}")

class Node:
    __slots__ = ("prior", "visits", "value_sum", "children")
    def __init__(self, prior):
        self.prior, self.visits, self.value_sum, self.children = prior, 0, 0.0, None
    @property
    def mean_value(self):
        return self.value_sum / self.visits if self.visits else 0.0


class MCTS:
    def __init__(self, prior_fn=None, value_fn=None, c_puct=1.4, seed=0, root_noise=None):
        self.prior_fn, self.value_fn, self.c_puct = prior_fn, value_fn, c_puct
        self.rng = random.Random(seed)
        # Exploration noise belongs at the root and nowhere else. Wrapping it
        # around prior_fn instead applies it at every expansion, which corrupts
        # the estimates the search is being run to produce.
        self.root_noise = root_noise

    def _priors(self, state):
        moves = state.agent_moves()
        if self.prior_fn is None:
            return {m: 1.0 / len(moves) for m in moves}
        raw = self.prior_fn(state)
        total = sum(raw.get(m, 0.0) for m in moves)
        if total <= 0:
            return {m: 1.0 / len(moves) for m in moves}
        return {m: raw.get(m, 0.0) / total for m in moves}

    def _evaluate(self, state):
        if self.value_fn is None:
            return playout(state, self.rng).outcome_for(state.to_play)
        return self.value_fn(state)

    def search(self, root_state, simulations):
        root = Node(1.0)
        priors = self._priors(root_state)
        if self.root_noise is not None:
            priors = self.root_noise(priors, self.rng)
        root.children = {m: Node(p) for m, p in priors.items()}
        for _ in range(simulations):
            node, state, path = root, root_state, [root]
            while node.children is not None and not state.is_over():
                total = math.sqrt(node.visits) or 1e-8
                move, child = max(
                    node.children.items(),
                    key=lambda kv: (-kv[1].mean_value if kv[1].visits else 0.0)
                    + self.c_puct * kv[1].prior * total / (1 + kv[1].visits))
                state = state.play(move); node = child; path.append(node)
            if state.is_over():
                value = state.outcome_for(state.to_play)
            else:
                node.children = {m: Node(p) for m, p in self._priors(state).items()}
                value = self._evaluate(state)
            for n in reversed(path):
                n.visits += 1; n.value_sum += value; value = -value
        return {m: c.visits for m, c in root.children.items()}

    def best_move(self, state, simulations):
        return max(self.search(state, simulations).items(), key=lambda kv: kv[1])[0]

centre = (SIZE // 2) * SIZE + SIZE // 2
for sims in (200, 2000):
    visits = MCTS(seed=0).search(Go(), sims)
    order = sorted(visits.items(), key=lambda kv: -kv[1])
    rank = [m for m, _ in order].index(centre) + 1
    print(f"{sims:>5} simulations: favourite = point {order[0][0]}, "
          f"centre ranked {rank} of {len(order)}")

def visits_to_probs(visits, temperature=1.0):
    if temperature <= 1e-6:
        best = max(visits.values())
        winners = [m for m, v in visits.items() if v == best]
        return {m: (1.0 / len(winners) if m in winners else 0.0) for m in visits}
    w = {m: v ** (1.0 / temperature) for m, v in visits.items()}
    total = sum(w.values()) or 1.0
    return {m: x / total for m, x in w.items()}


def generate_games(n_games, sims, seed=0, temp_moves=8, max_moves=200, prior_fn=None,
                   value_fn=None, root_noise=None, progress=None):
    """Play `n_games` self-play games and return per-position training records."""
    rng = random.Random(seed)
    boards, players, policies, values, move_nums, game_ids = [], [], [], [], [], []
    lengths, winners = [], []
    for g in range(n_games):
        mcts = MCTS(prior_fn=prior_fn, value_fn=value_fn, seed=seed * 1000 + g,
                    root_noise=root_noise)
        state, records = Go(), []
        while not state.is_over() and state.move_count < max_moves:
            visits = mcts.search(state, sims)
            if not visits:
                break
            total = sum(visits.values())
            policy = np.zeros(N_MOVES, dtype=np.float32)
            for m, v in visits.items():
                policy[N_POINTS if m == PASS else m] = v / total
            records.append((bytes(state.points), state.to_play, policy, state.move_count))
            if state.move_count < temp_moves:
                probs = visits_to_probs(visits, 1.0)
                move = rng.choices(list(probs), weights=list(probs.values()))[0]
            else:
                move = max(visits.items(), key=lambda kv: kv[1])[0]
            state = state.play(move)
        winner = state.winner()
        for board, player, policy, mv in records:
            boards.append(np.frombuffer(board, dtype=np.uint8))
            players.append(player); policies.append(policy)
            values.append(np.float32(1.0 if player == winner else -1.0))
            move_nums.append(mv); game_ids.append(g)
        lengths.append(state.move_count); winners.append(winner)
        if progress and (g + 1) % progress == 0:
            print(f"  {g+1}/{n_games} games", flush=True)
    return {
        "boards": np.stack(boards), "to_play": np.array(players, dtype=np.uint8),
        "policy": np.stack(policies), "value": np.array(values, dtype=np.float32),
        "move_number": np.array(move_nums, dtype=np.int16),
        "game_id": np.array(game_ids, dtype=np.int32),
        "lengths": lengths, "winners": winners,
    }

TEACHERS = [2, 10, 50, 200]
GAMES = 20_000

datasets = {}
for sims in TEACHERS:
    path = ART / f"dataset-5x5-{GAMES}g-{sims}s.npz"
    datasets[sims] = cached(
        path,
        lambda sims=sims: generate_games(GAMES, sims, seed=1 + sims, progress=1000),
        label=f"teacher with {sims} simulations",
    )

for sims, d in datasets.items():
    pol = d["policy"]
    floor = float(-(pol * np.log(np.clip(pol, 1e-9, None))).sum(1).mean())
    print(f"teacher {sims:>3} sims: {len(d['boards']):>7,} positions   "
          f"target entropy {floor:.2f} nats")

class GoNet(nn.Module):
    def __init__(self, channels=64, blocks=3, in_channels=3, groups=8):
        super().__init__()
        layers, prev = [], in_channels
        for _ in range(blocks):
            layers += [nn.Conv2d(prev, channels, 3, padding=1),
                       nn.GroupNorm(min(groups, channels), channels), nn.ReLU()]
            prev = channels
        self.trunk = nn.Sequential(*layers)
        self.policy_conv = nn.Conv2d(channels, 2, 1)
        self.policy_fc = nn.Linear(2 * N_POINTS, N_MOVES)
        self.value_conv = nn.Conv2d(channels, 1, 1)
        self.value_fc1 = nn.Linear(N_POINTS, 64)
        self.value_fc2 = nn.Linear(64, 1)

    def forward(self, x):
        h = self.trunk(x)
        p = F.relu(self.policy_conv(h)).flatten(1)
        v = F.relu(self.value_conv(h)).flatten(1)
        v = F.relu(self.value_fc1(v))
        return self.policy_fc(p), torch.tanh(self.value_fc2(v)).squeeze(-1)


def encode(boards, to_play):
    """(N, 25) boards -> (N, 3, 5, 5) planes."""
    grid = boards.reshape(-1, SIZE, SIZE)
    x = np.zeros((len(boards), 3, SIZE, SIZE), dtype=np.float32)
    x[:, 0] = (grid == BLACK)
    x[:, 1] = (grid == WHITE)
    x[:, 2] = (to_play == BLACK)[:, None, None]
    return x


def loss_fn(logits, value_pred, policy_target, value_target, value_weight=1.0):
    """Soft cross-entropy against the search's visit distribution, plus value MSE.

    The policy target is a *distribution*, not a label — it is what MCTS decided
    after spending its simulations — so this is the full soft cross-entropy
    rather than a classification loss.
    """
    logp = F.log_softmax(logits, dim=1)
    policy_loss = -(policy_target * logp).sum(1).mean()
    value_loss = F.mse_loss(value_pred, value_target)
    return policy_loss + value_weight * value_loss, policy_loss, value_loss

print(f"{sum(p.numel() for p in GoNet().parameters()):,} parameters")

def symmetry(x, policy, k):
    """One of the eight dihedral images; the pass logit is invariant."""
    rot, flip = divmod(k, 2)
    grid = policy[:, :N_POINTS].reshape(-1, 1, SIZE, SIZE)
    if rot:
        x = torch.rot90(x, rot, dims=(2, 3)); grid = torch.rot90(grid, rot, dims=(2, 3))
    if flip:
        x = torch.flip(x, dims=(3,)); grid = torch.flip(grid, dims=(3,))
    out = torch.empty_like(policy)
    out[:, :N_POINTS] = grid.reshape(-1, N_POINTS)
    out[:, N_POINTS] = policy[:, N_POINTS]
    return x.contiguous(), out


def train(model, data, epochs=20, batch=512, lr=2e-3, val_frac=0.1, log=False):
    """Train, and keep the *best* epoch rather than the last one.

    The split is by **game**, not by position. Positions from one game share an
    outcome and overlap heavily, so splitting positions at random puts near-copies
    of the same board — with the same value label — on both sides of the divide,
    and validation reports a number better than the model deserves.
    """
    n = len(data["boards"])
    rng = np.random.default_rng(0)
    if "game_id" in data:
        games = np.unique(data["game_id"])
        rng.shuffle(games)
        val_games = set(games[: max(1, int(len(games) * val_frac))].tolist())
        is_val = np.array([g in val_games for g in data["game_id"]])
        tr, va = np.flatnonzero(~is_val), np.flatnonzero(is_val)
    else:
        idx = rng.permutation(n)
        cut = int(n * (1 - val_frac))
        tr, va = idx[:cut], idx[cut:]

    def to_dev(sel):
        return (torch.from_numpy(encode(data["boards"][sel], data["to_play"][sel])).to(DEVICE),
                torch.from_numpy(data["policy"][sel]).to(DEVICE),
                torch.from_numpy(data["value"][sel]).to(DEVICE))

    xt, pt, vt = to_dev(tr)
    xv, pv, vv = to_dev(va)
    model = model.to(DEVICE)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs, eta_min=lr * 0.05)

    best_ce, best_state, history = float("inf"), None, []
    for epoch in range(epochs):
        model.train()
        for s in range(0, len(tr), batch):
            sel = torch.randperm(len(tr), device=DEVICE)[:batch]
            xb, pb = symmetry(xt[sel], pt[sel], int(torch.randint(8, (1,))))
            loss, _, _ = loss_fn(*model(xb), pb, vt[sel])
            opt.zero_grad(set_to_none=True); loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
        sched.step()
        model.eval()
        with torch.no_grad():
            logits, value = model(xv)
            _, ce, _ = loss_fn(logits, value, pv, vv)
            top1 = (logits.argmax(1) == pv.argmax(1)).float().mean().item()
            acc = ((value > 0) == (vv > 0)).float().mean().item()
        history.append((epoch + 1, ce.item(), top1, acc))
        if ce.item() < best_ce:
            best_ce = ce.item()
            best_state = {k: v.detach().clone() for k, v in model.state_dict().items()}
        if log:
            print(f"  epoch {epoch+1:>2}: CE {ce:.4f}  top1 {top1:.3f}  value-acc {acc:.3f}")
    model.load_state_dict(best_state)
    return model, history

students = {}
for sims, data in datasets.items():
    path = CKPT / f"student-teacher-{sims}s.pt"
    model = GoNet()
    if path.exists():
        print(f"loading {path}")
        model.load_state_dict(torch.load(path, map_location="cpu",
                                         weights_only=False)["state_dict"])
        students[sims] = model.to(DEVICE)
        continue
    t0 = time.time()
    model, hist = train(model, data, epochs=25)
    students[sims] = model
    print(f"teacher {sims:>3} sims -> CE {hist[-1][1]:.3f}  top1 {hist[-1][2]:.3f}  "
          f"value-acc {hist[-1][3]:.3f}  [{time.time()-t0:.0f}s]")

class NetPlayer:
    """Wraps a network; caches one forward pass per position."""
    def __init__(self, model, temperature=0.0, seed=0):
        self.model = model.eval().to("cpu")
        self.cache, self.rng, self.temperature = {}, random.Random(seed), temperature

    def _eval(self, state):
        key = (bytes(state.points), state.to_play)
        if key not in self.cache:
            x = torch.from_numpy(encode(np.frombuffer(bytes(state.points), dtype=np.uint8)[None],
                                        np.array([state.to_play], dtype=np.uint8)))
            with torch.no_grad():
                logits, value = self.model(x)
            self.cache[key] = (torch.softmax(logits, 1)[0].numpy(), float(value))
        return self.cache[key]

    def prior(self, state):
        probs = self._eval(state)[0]
        out = {m: float(probs[N_POINTS if m == PASS else m]) for m in state.agent_moves()}
        total = sum(out.values())
        return {m: v / total for m, v in out.items()} if total > 0 else \
               {m: 1 / len(out) for m in out}

    def value(self, state):
        return self._eval(state)[1]

    def select(self, state):
        prior = self.prior(state)
        if self.temperature <= 1e-6:
            return max(prior.items(), key=lambda kv: kv[1])[0]
        moves = list(prior)
        w = [prior[m] ** (1 / self.temperature) for m in moves]
        return self.rng.choices(moves, weights=w)[0]


class SearchPlayer:
    """MCTS with whichever of the two heads it is allowed to use."""
    def __init__(self, sims, prior_fn=None, value_fn=None, seed=0):
        self.sims = sims
        self.mcts = MCTS(prior_fn=prior_fn, value_fn=value_fn, seed=seed)

    def select(self, state):
        return self.mcts.best_move(state, self.sims)


class RandomPlayer:
    def __init__(self, seed=0):
        self.rng = random.Random(seed)
    def select(self, state):
        return rollout_move(state, self.rng)

def random_opening(seed, plies=2):
    rng = random.Random(seed)
    state = Go()
    for _ in range(plies):
        moves = state.agent_moves()
        if not moves:
            break
        state = state.play(rng.choice(moves))
    return state


def play_from(start, black, white, max_moves=200):
    state, players, n = start, {BLACK: black, WHITE: white}, 0
    while not state.is_over() and n < max_moves:
        state = state.play(players[state.to_play].select(state)); n += 1
    return state


def match(a, b, games=30, seed=0):
    """Games won by `a`, colours alternated, openings varied."""
    wins = 0
    for i in range(games):
        opening = random_opening(seed * 977 + i // 2)
        if i % 2 == 0:
            wins += play_from(opening, a, b).winner() == BLACK
        else:
            wins += play_from(opening, b, a).winner() == WHITE
    return wins

GAMES_PER_MATCH = 30
rows = []
for sims, model in students.items():
    student = NetPlayer(model)
    vs_random = match(student, RandomPlayer(seed=5), GAMES_PER_MATCH, seed=1)
    vs_teacher = match(student, SearchPlayer(sims, seed=7), GAMES_PER_MATCH, seed=2)
    rows.append((sims, vs_random, vs_teacher))
    print(f"student of the {sims:>3}-sim teacher: "
          f"vs random {vs_random:>2}/{GAMES_PER_MATCH}, "
          f"vs its own teacher {vs_teacher:>2}/{GAMES_PER_MATCH}")

sims_key = max(students)
model = students[sims_key]
net = NetPlayer(model)
baseline = NetPlayer(model)          # the same weights, no search

for budget in (10, 50):
    player = SearchPlayer(budget, prior_fn=net.prior, value_fn=net.value, seed=3)
    w = match(player, baseline, GAMES_PER_MATCH, seed=4)
    print(f"policy+value, {budget:>3} simulations vs no search at all: "
          f"{w:>2}/{GAMES_PER_MATCH}")

configs = {
    "uniform + rollout": (None, None),
    "policy  + rollout": (net.prior, None),
    "uniform + value  ": (None, net.value),
    "policy  + value  ": (net.prior, net.value),
}
BUDGET = 50
for name, (pf, vf) in configs.items():
    player = SearchPlayer(BUDGET, prior_fn=pf, value_fn=vf, seed=6)
    w = match(player, baseline, GAMES_PER_MATCH, seed=8)
    print(f"{name} @ {BUDGET} sims: {w:>2}/{GAMES_PER_MATCH} = {w/GAMES_PER_MATCH:.2f}")

def dirichlet(prior, rng, alpha=0.6, weight=0.25):
    """Noise on the root prior only, so every game can deviate somewhere."""
    if weight <= 0 or len(prior) < 2:
        return prior
    draws = [rng.gammavariate(alpha, 1.0) for _ in prior]
    total = sum(draws) or 1.0
    return {m: (1 - weight) * p + weight * (d / total)
            for (m, p), d in zip(prior.items(), draws)}


def selfplay_generation(model, games, sims, seed, anchor=None, anchor_fraction=0.5,
                        buffer=None, window=10, epochs=6):
    """One turn of the crank: play, add to the buffer, train, return the next network.

    ``model=None`` means no network exists yet, so the games are played by
    classical MCTS — uniform prior, random rollouts. That is what this loop is
    before its first network, and it is how a cold start begins.
    """
    noise_rng = random.Random(seed)
    if model is None:
        fresh = generate_games(games, sims, seed=seed)          # classical MCTS
    else:
        player = NetPlayer(model)
        fresh = generate_games(
            games, sims, seed=seed,
            prior_fn=player.prior, value_fn=player.value,
            root_noise=lambda priors, rng: dirichlet(priors, rng),
        )
    buffer = (buffer or []) + [fresh]
    buffer = buffer[-window:]
    # `game_id` has to survive the merge, or `train` silently falls back to a
    # position-level split — the leak the by-game split exists to prevent. Ids are
    # offset per generation so two generations never collide on the same number.
    merged = {k: np.concatenate([b[k] for b in buffer])
              for k in ("boards", "to_play", "policy", "value")}
    merged["game_id"] = np.concatenate(
        [b["game_id"] + i * 1_000_000 for i, b in enumerate(buffer)])

    # Mixing the supervised games back in is not a nicety. Trained only on the
    # fresh buffer — a thousand games against the twenty thousand the network
    # already learned from — it forgets, and strength falls monotonically with
    # every generation.
    if anchor is not None:
        n_self = len(merged["boards"])
        n_anchor = int(n_self * anchor_fraction / max(1e-9, 1 - anchor_fraction))
        take = np.random.default_rng(seed).choice(len(anchor["boards"]),
                                                  size=min(n_anchor, len(anchor["boards"])),
                                                  replace=False)
        anchor_slice = {k: anchor[k][take] for k in ("boards", "to_play", "policy", "value")}
        anchor_slice["game_id"] = anchor["game_id"][take] + 900_000_000
        merged = {k: np.concatenate([merged[k], anchor_slice[k]]) for k in merged}

    nxt, hist = train(GoNet(), merged, epochs=epochs)
    return nxt, buffer, fresh, hist

GENERATIONS = 40
SELFPLAY_GAMES = 300
SELFPLAY_SIMS = 100
SELFPLAY_DIR = Path("scripts/fixtures/selfplay")

generations = []
if (SELFPLAY_DIR / f"gen-{GENERATIONS:03d}.pt").exists():
    for gen in range(1, GENERATIONS + 1):
        blob = torch.load(SELFPLAY_DIR / f"gen-{gen:03d}.pt", map_location="cpu",
                          weights_only=False)
        m = GoNet(); m.load_state_dict(blob["state_dict"]); generations.append(m)
    print(f"loaded {GENERATIONS} generations from {SELFPLAY_DIR}")
else:
    SELFPLAY_DIR.mkdir(parents=True, exist_ok=True)
    current, buffer = None, []          # None means "play classically this round"
    for gen in range(1, GENERATIONS + 1):
        t0 = time.time()
        current, buffer, fresh, hist = selfplay_generation(
            current, SELFPLAY_GAMES, SELFPLAY_SIMS, seed=100 + gen, buffer=buffer,
            window=12, epochs=8)
        generations.append(current)
        # Save what the cells below read: the games (the opening table walks
        # them) and the checkpoint (so the next run loads instead of recomputing).
        np.savez_compressed(SELFPLAY_DIR / f"games-gen-{gen:03d}.npz",
                            **{k: v for k, v in fresh.items()
                               if isinstance(v, np.ndarray)})
        torch.save({"state_dict": current.state_dict()},
                   SELFPLAY_DIR / f"gen-{gen:03d}.pt")
        distinct = len({bytes(b) for b in fresh["boards"]}) / len(fresh["boards"])
        print(f"gen {gen}: {distinct:.0%} distinct positions, "
              f"{np.mean(fresh['lengths']):.1f} moves, "
              f"pass mass {fresh['policy'][:, N_POINTS].mean():.3f}, "
              f"CE {hist[-1][1]:.3f} [{time.time()-t0:.0f}s]")

CENTRE = (SIZE // 2) * SIZE + SIZE // 2
games_files = sorted(SELFPLAY_DIR.glob("games-gen-*.npz"))
print(f"{'gen':>4}  {'favourite':>10}  {'share':>7}  {'entropy':>8}")
for f in games_files[:3] + games_files[4::6]:
    gen = int(f.stem.split("-")[-1])
    d = np.load(f)
    opening = d["policy"][d["move_number"] == 0]
    counts = np.bincount([int(p.argmax()) for p in opening], minlength=N_MOVES)
    top = int(counts.argmax())
    mean_pol = opening.mean(0)
    ent = float(-(mean_pol * np.log(np.clip(mean_pol, 1e-9, None))).sum())
    lbl = "pass" if top == N_POINTS else "ABCDE"[top % SIZE] + str(SIZE - top // SIZE)
    print(f"{gen:>4}  {lbl:>10}  {counts[top]/len(opening):>6.0%}  {ent:>8.2f}")
print(f"\nuniform entropy would be {math.log(N_MOVES):.2f}; "
      f"the solved optimum is the centre, {'ABCDE'[CENTRE % SIZE]}{SIZE - CENTRE // SIZE}")

sample = [1, 5, 10, 20, 40]
field = [("rollout-MCTS", SearchPlayer(200, seed=11)),
         ("supervised", NetPlayer(students[200]))]
field += [(f"gen-{i}", NetPlayer(generations[i - 1]))
          for i in sample if 0 < i <= len(generations)]

GAMES_PER_PAIR = 20
print(f"{'':>14}" + "".join(f"{n:>14}" for n, _ in field))
cache = {}
for i, (name_a, a) in enumerate(field):
    cells = []
    for j, (name_b, b) in enumerate(field):
        if i == j:
            cells.append("—"); continue
        key = (min(i, j), max(i, j))
        if key not in cache:
            cache[key] = match(field[key[0]][1], field[key[1]][1],
                               GAMES_PER_PAIR, seed=20 + key[0] * 7 + key[1])
        w = cache[key] if i < j else GAMES_PER_PAIR - cache[key]
        cells.append(f"{w}/{GAMES_PER_PAIR}")
    print(f"{name_a:>14}" + "".join(f"{c:>14}" for c in cells))