#!/usr/bin/env -S uv run --script
#
# /// script
# requires-python = ">=3.14"
# dependencies = [
#   "numpy>=2.5.3",
#   "torch>=2.14.0",
#   "tqdm>=4.70.0",
# ]
#
# [tool.uv.sources]
# torch = { index = "pytorch-cpu" }
#
# [[tool.uv.index]]
# name = "pytorch-cpu"
# url = "https://download.pytorch.org/whl/cpu"
# explicit = true
# ///

#
# This file is part of ctak.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with ctak. If not, see <https://www.gnu.org/licenses/>.

import argparse
import concurrent.futures
import json
import math
import os
import random
import subprocess
from collections.abc import Iterator
from dataclasses import dataclass
from itertools import product
from pathlib import Path

import numpy as np
import torch
from torch import Tensor
from torch.nn import Embedding, Linear, Module, MSELoss
from tqdm import tqdm

# ===================================================================
# Net dimensions
# ===================================================================

# (input dim, hidden dim) per board size; must match include/weights_<size>.h
NET_DIMS: dict[int, tuple[int, int]] = {
    5: (27, 64),  # FIVE_INP_NUM, FIVE_DENSE_NUM
    6: (38, 64),  # SIX_INP_NUM, SIX_DENSE_NUM
}


# ===================================================================
# C weight formatting
# ===================================================================


def _fmt_flat(a: np.ndarray, per_line: int = 8) -> str:
    a = np.ravel(np.asarray(a, dtype=np.float32))
    lines: list[str] = []
    for i in range(0, len(a), per_line):
        chunk = ", ".join(f"{v:.9g}" for v in a[i : i + per_line])
        if i + per_line < len(a):
            chunk += ","
        lines.append("    " + chunk)
    return "\n".join(lines)


def _fmt_matrix(mat: np.ndarray, per_line: int = 8) -> str:
    mat = np.asarray(mat, dtype=np.float32)
    rows: list[str] = []
    for r in mat:
        chunks = [
            ", ".join(f"{v:.9g}" for v in r[i : i + per_line])
            for i in range(0, len(r), per_line)
        ]
        rows.append("    {" + ",\n     ".join(chunks) + "},")
    return "\n".join(rows)


# ===================================================================
# Config
# ===================================================================


@dataclass(frozen=True)
class DataGenConfig:
    size: int
    depth: int
    min_ply: int
    games: int
    procs: int
    replay_iters: int
    root: Path
    train_only_self_play: bool
    k: float


@dataclass(frozen=True)
class TrainingGridSearchConfig:
    smooth_factor: list[float]
    epochs: list[int]
    batch: list[int]
    lr: list[float]
    val_frac: list[float]

    def make_configs(self, data_gen_config: DataGenConfig) -> list["Candidate"]:
        inp_dim, hidden = NET_DIMS[data_gen_config.size]
        flds = self.__dataclass_fields__.keys()
        almost_cands = [
            dict(zip(flds, vals)) for vals in product(*(getattr(self, k) for k in flds))
        ]
        for ident, almost_cand in enumerate(almost_cands):
            net = Net(inp_dim, hidden)
            net.to_c(
                size=data_gen_config.size,
                ident=ident,
                path=data_gen_config.root
                / "include"
                / f"cand_{ident}_weights_{data_gen_config.size}.c",
                note=f"{ident=}, v000, random init",
            )
            cand_copy_for(
                6 if data_gen_config.size == 5 else 5, ident, data_gen_config.root
            )
            almost_cand["net"] = net
            almost_cand["train_only_self_play"] = data_gen_config.train_only_self_play
        return [Candidate(**d) for d in almost_cands]


@dataclass
class Candidate:
    train_only_self_play: bool
    smooth_factor: float
    epochs: int
    batch: int
    lr: float
    val_frac: float
    net: "Net"
    elo: float = 1500.0
    rated_games: int = 0

    def __str__(self) -> str:
        return f"(λ={self.smooth_factor}, E={self.epochs}, BS={self.batch}, η={self.lr:.2e}, V%={self.val_frac * 100.0:.1f})"

    def training_pool(self, pool: Pool) -> Pool:
        if self.train_only_self_play:
            return pool.select_chains(pool.chain_self_play)
        return pool

    def fit(self, pool: Pool) -> tuple[float, float]:
        # The train/val split is fixed but targets are refreshed
        split = DataSplit.build(self.training_pool(pool), self.val_frac)
        if split.n == 0:
            return math.nan, math.nan

        optimizer = torch.optim.AdamW(self.net.parameters(), lr=self.lr)
        loss_fn = MSELoss()

        pbar = tqdm(range(self.epochs), desc="train", unit="ep")
        for _ in pbar:
            split.refresh_targets(self.net, self.smooth_factor)

            self.net.train()
            for head, cells, parity, target in split.batches(self.batch):
                optimizer.zero_grad()
                predictions = self.net(head, cells) * parity
                loss = loss_fn(predictions, target)
                loss.backward()
                optimizer.step()

            with torch.no_grad():
                self.net.eval()
                val = split.val_idx
                predictions = self.net(split.pool.head[val], split.pool.cells[val])
                predictions = predictions * split.pool.parity[val]
                targets = split.targets[val]
                mse = loss_fn(predictions, targets).item()
                corr = torch.corrcoef(torch.stack([predictions, targets]))[0, 1].item()

            pbar.set_postfix(val_corr=f"{corr:.4f}", val_mse=f"{mse:.4f}")

        return corr, mse


# ===================================================================
# Self-play records: games (for ratings) and chains (for training)
# ===================================================================


@dataclass(frozen=True)
class Games:
    black: np.ndarray  # [G] candidate id playing black
    white: np.ndarray  # [G] candidate id playing white
    result: np.ndarray  # [G] outcome from black's perspective, {-1, 0, 1}
    newest: np.ndarray  # [G] game comes from the last loaded file

    def head_to_head(self) -> Iterator[tuple[int, int, float]]:
        for i in np.flatnonzero(self.newest & (self.black != self.white)):
            black, white, result = (
                int(self.black[i]),
                int(self.white[i]),
                float(self.result[i]),
            )
            # elo score: 1 win, 0.5 draw, 0 loss, from a's perspective
            yield black, white, (result + 1.0) / 2.0


@dataclass(frozen=True)
class Pool:
    head: Tensor  # [n, 2] fixed head inputs of the root state s_t
    cells: Tensor  # [n, n_cells] cell states 0..6 of s_t
    leaf_head: Tensor  # [n, 2] head inputs of the PV leaf pv_t
    leaf_cells: Tensor  # [n, n_cells] cell states of pv_t
    parity: Tensor  # [n] +1 when black is to move at s_t
    leaf_parity: Tensor  # [n] +1 when black is to move at pv_t
    terminal: Tensor  # [n] leaf is an exact terminal
    terminal_value: Tensor  # [n] exact terminal value, black's perspective
    chain_lens: Tensor  # [C] rows per chain
    chain_result: Tensor  # [C] game result, black's perspective
    chain_self_play: Tensor  # [C] chain comes from a self vs self game

    @classmethod
    def empty(cls, n_cells: int) -> "Pool":
        def rows(dtype: torch.dtype, width: int = 0) -> Tensor:
            return torch.empty((0, width) if width else (0,), dtype=dtype)

        return cls(
            head=rows(torch.float32, 2),
            cells=rows(torch.uint8, n_cells),
            leaf_head=rows(torch.float32, 2),
            leaf_cells=rows(torch.uint8, n_cells),
            parity=rows(torch.float32),
            leaf_parity=rows(torch.float32),
            terminal=rows(torch.bool),
            terminal_value=rows(torch.float32),
            chain_lens=rows(torch.int64),
            chain_result=rows(torch.float32),
            chain_self_play=rows(torch.bool),
        )

    @property
    def n_rows(self) -> int:
        return self.head.shape[0]

    def select_chains(self, keep: Tensor) -> "Pool":
        keep = keep.bool()
        rows = torch.repeat_interleave(
            torch.arange(self.chain_lens.shape[0], dtype=torch.int64), self.chain_lens
        )
        idx = torch.nonzero(keep[rows]).flatten()
        return Pool(
            head=self.head[idx],
            cells=self.cells[idx],
            leaf_head=self.leaf_head[idx],
            leaf_cells=self.leaf_cells[idx],
            parity=self.parity[idx],
            leaf_parity=self.leaf_parity[idx],
            terminal=self.terminal[idx],
            terminal_value=self.terminal_value[idx],
            chain_lens=self.chain_lens[keep],
            chain_result=self.chain_result[keep],
            chain_self_play=self.chain_self_play[keep],
        )


# ===================================================================
# Reading game records
# ===================================================================

_ROW_FIELDS = (
    "head",
    "cells",
    "leaf_head",
    "leaf_cells",
    "parity",
    "leaf_parity",
    "terminal",
    "terminal_value",
)


def _chain_arrays(
    rows: list[dict], n_cells: int, source: Path
) -> dict[str, np.ndarray]:

    def column(part: str, key: str, dtype: np.dtype) -> np.ndarray:
        return np.array([r[part][key] for r in rows], dtype=dtype)

    out = {
        "head": column("root", "head", np.float32),
        "cells": column("root", "cells", np.uint8),
        "leaf_head": column("leaf", "head", np.float32),
        "leaf_cells": column("leaf", "cells", np.uint8),
        "parity": column("root", "parity", np.float32),
        "leaf_parity": column("leaf", "parity", np.float32),
        "terminal": column("leaf", "terminal", np.bool_),
        "terminal_value": column("leaf", "terminal_value", np.float32),
    }
    for name in ("cells", "leaf_cells"):
        if out[name].shape[1] != n_cells:
            raise ValueError(
                f"{source}: {name} width {out[name].shape[1]}, expected {n_cells}"
            )
    return out


def _game_colours(game: dict) -> tuple[int, int]:
    a, b = int(game["cand_A"]), int(game["cand_B"])
    for r in game["rows"]:
        if r["root"]["parity"] > 0:
            cand = int(r["cand"])
            return cand, b if cand == a else a
    recorded = {int(r["cand"]) for r in game["rows"]}
    return (a, b) if a not in recorded else (b, a)


def load_window(
    files: list[Path], size: int, n_candidates: int
) -> tuple[list[Pool], Games]:
    n_cells = NET_DIMS[size][0] - 2
    blocks = [{f: [] for f in _ROW_FIELDS} for _ in range(n_candidates)]
    chain_lens: list[list[int]] = [[] for _ in range(n_candidates)]
    chain_result: list[list[float]] = [[] for _ in range(n_candidates)]
    chain_self_play: list[list[bool]] = [[] for _ in range(n_candidates)]
    black, white, result, newest = [], [], [], []

    for file_index, path in enumerate(files):
        is_newest = file_index == len(files) - 1
        with open(path) as f:
            for line in f:
                if not line.strip():
                    continue
                game = json.loads(line)
                game_result = float(game["result"])
                b, w = _game_colours(game)
                black.append(b)
                white.append(w)
                result.append(game_result)
                newest.append(is_newest)

                by_cand: dict[int, list[dict]] = {}
                for r in game["rows"]:
                    by_cand.setdefault(int(r["cand"]), []).append(r)

                for cand, rows in by_cand.items():
                    parity = [r["root"]["parity"] for r in rows]
                    for sign in (1, -1):
                        chain = [
                            r for r, p in zip(rows, parity) if (p > 0) == (sign > 0)
                        ]
                        if not chain:
                            continue
                        for name, array in _chain_arrays(chain, n_cells, path).items():
                            blocks[cand][name].append(array)
                        chain_lens[cand].append(len(chain))
                        chain_result[cand].append(game_result)
                        chain_self_play[cand].append(b == w)

    def pool_of(cand: int) -> Pool:
        if not chain_lens[cand]:
            return Pool.empty(n_cells)
        arrays = {
            name: torch.from_numpy(np.concatenate(blocks[cand][name]))
            for name in _ROW_FIELDS
        }
        return Pool(
            **arrays,
            chain_lens=torch.tensor(chain_lens[cand], dtype=torch.int64),
            chain_result=torch.tensor(chain_result[cand], dtype=torch.float32),
            chain_self_play=torch.tensor(chain_self_play[cand], dtype=torch.bool),
        )

    games = Games(
        black=np.array(black, dtype=np.int16),
        white=np.array(white, dtype=np.int16),
        result=np.array(result, dtype=np.float32),
        newest=np.array(newest, dtype=np.bool_),
    )
    return [pool_of(c) for c in range(n_candidates)], games


# ===================================================================
# Training data: fixed rows, refreshed targets
# ===================================================================


@dataclass
class DataSplit:
    pool: Pool
    train_idx: Tensor
    val_idx: Tensor
    targets: Tensor  # [n] aligned with pool rows, refreshed per epoch
    chain_end: Tensor  # [C] last row index of each chain
    wave_order: Tensor  # [n] rows ordered by distance from their chain's end
    wave_bounds: Tensor  # [d] start of distance d in wave_order

    @property
    def n(self) -> int:
        return self.pool.n_rows

    @classmethod
    def build(cls, pool: Pool, val_frac: float) -> "DataSplit":
        n = pool.n_rows
        if n == 0:
            empty = torch.empty(0, dtype=torch.int64)
            return cls(pool, empty, empty, torch.empty(0), empty, empty, empty)

        # Smoothing is G[t] = (1-λ)V[t] + λG[t+1] with G at the chain end.
        # Chains are contiguous, so every row that is d plies from its chain's
        # end can be updated in one vectorised wavefront step per d.
        lens = pool.chain_lens
        chain_end = torch.cumsum(lens, dim=0) - 1
        chain_of_row = torch.repeat_interleave(
            torch.arange(lens.shape[0], dtype=torch.int64), lens
        )
        dist = chain_end[chain_of_row] - torch.arange(n, dtype=torch.int64)
        wave_order = torch.argsort(dist, stable=True)
        wave_bounds = torch.searchsorted(
            dist[wave_order], torch.arange(int(dist.max()) + 2, dtype=torch.int64)
        )

        n_val = max(1, int(n * val_frac))
        perm = torch.randperm(n)
        return cls(
            pool=pool,
            train_idx=perm[n_val:],
            val_idx=perm[:n_val],
            targets=torch.zeros(n),
            chain_end=chain_end,
            wave_order=wave_order,
            wave_bounds=wave_bounds,
        )

    def refresh_targets(self, net: Module, lam: float) -> None:
        pool = self.pool
        net.eval()
        with torch.no_grad():
            # Net(pv_t), black's perspective; exact terminals override the net
            leaf = net(pool.leaf_head, pool.leaf_cells) * pool.leaf_parity
            leaf = torch.where(pool.terminal, pool.terminal_value, leaf)

        targets = torch.zeros_like(leaf)
        targets[self.chain_end] = pool.chain_result
        for lo, hi in zip(self.wave_bounds[1:-1], self.wave_bounds[2:]):
            idx = self.wave_order[lo:hi]
            targets[idx] = (1.0 - lam) * leaf[idx] + lam * targets[idx + 1]
        self.targets = targets

    def batches(
        self, batch_size: int
    ) -> Iterator[tuple[Tensor, Tensor, Tensor, Tensor]]:
        order = torch.randperm(self.train_idx.shape[0])
        train_targets = self.targets[self.train_idx]
        for start in range(0, order.shape[0], batch_size):
            idx = order[start : start + batch_size]
            rows = self.train_idx[idx]
            yield (
                self.pool.head[rows],
                self.pool.cells[rows],
                self.pool.parity[rows],
                train_targets[idx],
            )


# ===================================================================
# Network
# ===================================================================


def fast_tanh(x: Tensor) -> Tensor:
    # Must match fast_tanh in include/nn.c: degree-11 fit of tanh on [-3, 3],
    # saturating to +-1 outside.
    x2 = x * x
    poly = x * (
        0.997136032
        + x2 * (
            -0.305215124
            + x2 * (0.0832466371 + x2 * (-0.0139947924 + x2 * (0.00123106314 + x2 * -4.29748644e-05)))
        )
    )
    return torch.where(x > 3.0, 1.0, torch.where(x < -3.0, -1.0, poly))


def fast_silu(x: Tensor) -> Tensor:
    # Must match fast_silu in include/nn.c: [4,2] Pade approximant with cutoff.
    # The cutoff below is evaluated at x = 0 (instead of x) because torch
    # evaluates both branches of where: 1 + 0.1 * x is 0 for x = -10.
    xs = torch.where(x < -4.6, 0.0, x)
    poly = xs * (0.5 + xs * (0.25 + xs * (0.05 + 0.004167 * xs))) / (1 + 0.1 * xs)
    return torch.where(x < -4.6, 0.0, torch.where(x > 4.6, x, poly))


class Net(Module):
    def __init__(self, inp_dim: int, hidden: int) -> None:
        super().__init__()
        # learned cell-state embedding, see nn.c
        self.embedding = Embedding(7, 1)
        self.hidden = Linear(inp_dim, hidden, bias=True)
        self.output = Linear(hidden, 1, bias=True)

    def forward(self, head: Tensor, cells: Tensor) -> Tensor:
        # head: [n, 2] fixed head inputs; cells: [n, n_cells] raw cell states
        # (0..6) - the only place they become anything other than integers
        inp = torch.cat([head, self.embedding(cells.long()).flatten(1)], dim=1)
        hidden = fast_silu(self.hidden(inp))
        return fast_tanh(self.output(hidden)).squeeze(-1)

    def to_c(
        self,
        *,
        size: int,
        ident: int | None,
        path: Path,
        note: str,
    ) -> None:
        upper = "FIVE" if size == 5 else "SIX"
        prefix = (f"cand_{ident}_" if ident is not None else "") + upper
        guard = f"WEIGHTS_{upper}" + (f"_{ident}" if ident is not None else "")

        E = self.embedding.weight.detach().cpu().numpy()
        W = self.hidden.weight.detach().cpu().numpy()
        B = self.hidden.bias.detach().cpu().numpy()
        LW = self.output.weight.detach().cpu().numpy().ravel()
        LB = float(self.output.bias.detach().cpu().item())

        with open(path, "w") as f:
            lines = [
                "/* generated by nn_train.py, do not edit by hand */\n",
                f"/* {note} */\n",
                f"#ifndef {guard}\n#define {guard}\n",
                f'#include "weights_{size}.h"\n\n',
                f"const float {prefix}_embedding[{len(E)}] =\n{{\n{_fmt_flat(E, per_line=7)}\n}};\n\n",
                f"const float {prefix}_dense_weights[{upper}_DENSE_NUM][{upper}_INP_NUM] NN_ALIGNED =\n{{\n{_fmt_matrix(W)}\n}};\n\n",
                f"const float {prefix}_dense_biases[{upper}_DENSE_NUM] =\n{{\n{_fmt_flat(B)}\n}};\n\n",
                f"const float {prefix}_logit_weights[{upper}_DENSE_NUM] NN_ALIGNED =\n{{\n{_fmt_flat(LW)}\n}};\n\n",
                f"const float {prefix}_logit_bias = {LB:.9g};\n",
                "#endif",
            ]
            for line in lines:
                _ = f.write(line)


# ===================================================================
# Process helpers
# ===================================================================


def fan_out(
    cmds: list[list[str]], n: int, procs: int, tag: str, root: Path
) -> Iterator[str]:
    with concurrent.futures.ThreadPoolExecutor(max_workers=procs) as pool:
        futs = [
            pool.submit(
                subprocess.run,
                cmd,
                cwd=root,
                check=True,
                capture_output=True,
                text=True,
            )
            for _ in range(n)
            for cmd in cmds
        ]
        for fut in tqdm(
            concurrent.futures.as_completed(futs),
            total=n * len(cmds),
            desc=tag,
            unit="game",
        ):
            yield fut.result().stdout


def cand_copy_for(size: int, id: int, root: Path) -> None:
    src = root / "include" / f"weights_{size}.c"
    if not src.is_file():
        raise SystemExit(f"{src} not found; build a size-{size} champion first")
    upper = "FIVE" if size == 5 else "SIX"
    txt = src.read_text()
    for sym in (
        "embedding",
        "dense_weights",
        "dense_biases",
        "logit_weights",
        "logit_bias",
    ):
        txt = txt.replace(f"{upper}_{sym}", f"cand_{id}_{upper}_{sym}")
    txt = txt.replace(f"WEIGHTS_{upper}", f"WEIGHTS_{upper}_{id}")
    _ = (root / "include" / f"cand_{id}_weights_{size}.c").write_text(txt)


# ===================================================================
# The bootstrap loop
# ===================================================================


def update_elo(a: Candidate, b: Candidate, score: float, k: float) -> None:
    # score is a's result as an elo score: 1.0 win, 0.5 draw, 0.0 loss
    ea = 1.0 / (1.0 + 10.0 ** ((b.elo - a.elo) / 400.0))
    delta = k * (score - ea)
    a.elo += delta
    b.elo -= delta


def run_step(
    *,
    candidates: list[Candidate],
    data_gen_config: DataGenConfig,
    iter: int,
):
    def data_path(i: int) -> Path:
        return (
            data_gen_config.root
            / "data"
            / f"selfplay-{data_gen_config.size}-v{i:03d}.json"
        )

    print(f"\n== iteration {iter} ==")

    num_candidates = len(candidates)
    cand_pairs = [
        (i, j) for i in range(num_candidates) for j in range(i, num_candidates)
    ]

    # 1. rebuild
    for i, j in cand_pairs:
        _ = subprocess.run(
            ["make", f"selfplay-{i}-{j}"],
            cwd=data_gen_config.root,
            check=True,
            capture_output=True,
        )

    # 2. generate self-play data
    data = data_path(iter)
    print(
        f"selfplay: {num_candidates=}, {data_gen_config.games=}, "
        f"{data_gen_config.size=}, {data_gen_config.depth=}, {data_gen_config.min_ply=}, "
        f"{data_gen_config.train_only_self_play=}"
    )
    out_games: list[str] = []
    for out in fan_out(
        [
            [
                f"./selfplay-{i}-{j}",
                str(data_gen_config.size),
                str(data_gen_config.depth),
                str(data_gen_config.min_ply),
            ]
            for (i, j) in cand_pairs
        ],
        data_gen_config.games,
        data_gen_config.procs,
        "selfplay",
        data_gen_config.root,
    ):
        out_games.extend(out.splitlines())
    print(f"writing {len(out_games):,} games to {data}")
    with open(data, "w") as f:
        _ = f.write("\n".join(out_games))
        if out_games:
            _ = f.write("\n")

    # 3. rate every head-to-head game played in this iteration
    data_files = [
        data_path(i)
        for i in range(max(iter - data_gen_config.replay_iters, 1), iter + 1)
    ]
    print("Loading data from", [str(c) for c in data_files], "...", end=" ")
    pools, games = load_window(data_files, data_gen_config.size, num_candidates)
    print("candidates have", [f"{p.n_rows:,}" for p in pools], "rows.")

    fresh = list(games.head_to_head())
    random.shuffle(fresh)
    for a, b, score in fresh:
        update_elo(candidates[a], candidates[b], score, data_gen_config.k)
        candidates[a].rated_games += 1
        candidates[b].rated_games += 1

    total_rated_games = sum(c.rated_games for c in candidates)
    print(
        f"Tournament standings ({total_rated_games:,} rated games, {total_rated_games / len(candidates):,} games/candidate):"
    )
    for idx, cand in sorted(
        enumerate(candidates), key=lambda t: t[1].elo, reverse=True
    ):
        print(f"\tcandidate {idx:03d}: Elo {cand.elo:4.1f} Config {cand}")

    # 4. write champ, unconditionally
    idx, champ = max(enumerate(candidates), key=lambda t: t[1].elo)
    path = data_gen_config.root / "include" / f"weights_{data_gen_config.size}.c"
    champ.net.to_c(
        size=data_gen_config.size,
        ident=None,
        path=path,
        note=f"{iter=}, elo={champ.elo:.1f}, rated_games={champ.rated_games:,}",
    )
    _ = subprocess.run(
        ["make", "clean"], cwd=data_gen_config.root, check=True, capture_output=True
    )
    _ = subprocess.run(
        ["make", "native"], cwd=data_gen_config.root, check=True, capture_output=True
    )
    print(f"Wrote candidate {idx:03d} to {path!s}")

    # 5. fine-tune on the hybrid targets -> candidate
    for idx, cand in enumerate(candidates):
        corr, mse = cand.fit(pools[idx])
        cand.net.to_c(
            size=data_gen_config.size,
            ident=idx,
            path=data_gen_config.root
            / "include"
            / f"cand_{idx}_weights_{data_gen_config.size}.c",
            note=f"id={idx}, v{iter:03d}, val_corr={corr:.4f}, val_mse={mse:.4f}",
        )


# ===================================================================
# Main
# ===================================================================


def main() -> None:
    root = Path(__file__).resolve().parent.parent
    p = argparse.ArgumentParser(
        description="Train the neural network evaluators of ctak",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    _ = p.add_argument(
        "--size",
        type=int,
        choices=sorted(NET_DIMS),
        default=5,
        help="board size",
    )
    _ = p.add_argument(
        "--iters",
        type=int,
        default=100,
        help="max loop iterations",
    )
    _ = p.add_argument(
        "--games",
        type=int,
        default=200,
        help="games per iteration",
    )
    _ = p.add_argument(
        "--replay-iters",
        type=int,
        default=10,
        help="grid points for number of previous data sources to retain",
    )
    _ = p.add_argument(
        "--smooth-factor",
        type=float,
        nargs="+",
        default=[0.3, 0.5, 0.7, 0.9],
        help="grid points for the EMA smoothing factor",
    )
    _ = p.add_argument(
        "--epochs",
        type=int,
        nargs="+",
        default=[5],
        help="grid points for candidate training epochs",
    )
    _ = p.add_argument(
        "--lr",
        type=float,
        nargs="+",
        default=[1e-4],
        help="grid points for learning rate",
    )
    _ = p.add_argument(
        "--batch",
        type=int,
        nargs="+",
        default=[128],
        help="grid points for training batch size",
    )
    _ = p.add_argument(
        "--val-frac",
        type=float,
        nargs="+",
        default=[0.05],
        help="grid points for fraction of rows held out for validation",
    )
    _ = p.add_argument(
        "--procs",
        type=int,
        default=None,
        help="parallel game processes (number of CPUs if unset)",
    )
    _ = p.add_argument(
        "--depth",
        type=int,
        default=4,
        help="search depth for selfplay and the matches",
    )
    _ = p.add_argument(
        "--min-ply",
        type=int,
        default=8,
        help="length of random opening in games",
    )
    _ = p.add_argument(
        "--train-all-games",
        action="store_true",
        help="train on all selfplay games, including matches between "
        "different candidates (default: self vs self games only)",
    )
    _ = p.add_argument(
        "--k-factor",
        type=float,
        default=32.0,
        help="elo k-factor",
    )

    args = p.parse_args()

    # avoid counter-play start
    if args.min_ply < 2:
        _ = p.error("--min-ply must be >= 2")

    if args.procs is None:
        args.procs = max(1, (os.cpu_count() or 1))

    data_gen = DataGenConfig(
        size=args.size,
        depth=args.depth,
        min_ply=args.min_ply,
        games=args.games,
        procs=args.procs,
        replay_iters=args.replay_iters,
        root=root,
        train_only_self_play=not args.train_all_games,
        k=args.k_factor,
    )
    tgsc = TrainingGridSearchConfig(
        smooth_factor=args.smooth_factor,
        epochs=args.epochs,
        batch=args.batch,
        lr=args.lr,
        val_frac=args.val_frac,
    )

    candidates = tgsc.make_configs(data_gen)

    (root / "data").mkdir(exist_ok=True)

    num_params = sum(
        p.numel() for p in candidates[0].net.parameters() if p.requires_grad
    )
    print(f"Number of trainable parameters: {num_params:,}")

    for iter in range(1, args.iters + 1):
        run_step(
            candidates=candidates,
            data_gen_config=data_gen,
            iter=iter,
        )
    print(f"\n== done, champion in include/weights_{args.size}.c ==")


if __name__ == "__main__":
    main()