#!/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/>.

from collections.abc import Iterator
from dataclasses import dataclass
from itertools import product
from math import log, ceil

import argparse
import concurrent.futures
import json
import os
from pathlib import Path
import subprocess

import numpy as np
import torch
from tqdm import tqdm

from torch.nn import Embedding, MSELoss, Linear, Module
from torch.nn.functional import silu, tanh
from torch import Tensor


# ===================================================================
# 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


@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) -> list[CandidateTrainingConfig]:
        flds = self.__dataclass_fields__.keys()
        return [
            CandidateTrainingConfig(**dict(zip(flds, vals)))
            for vals in product(*(getattr(self, k) for k in flds))
        ]


@dataclass(frozen=True)
class CandidateTrainingConfig:
    smooth_factor: float
    epochs: int
    batch: int
    lr: float
    val_frac: float

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


# ===================================================================
# Self-play records
# ===================================================================


def smooth(
    leaf_values: Tensor,
    result: float,
    lam: float,
) -> Tensor:
    # Smooth(leaf_values, result, lambda), rooted at the game outcome:
    # G[-1] = result, G[t] = (1-lambda)V[t] + lambda G[t+1]
    targets = torch.empty_like(leaf_values)
    targets[-1] = result

    for ply in range(len(leaf_values) - 2, -1, -1):
        targets[ply] = (1.0 - lam) * leaf_values[ply] + lam * targets[ply + 1]

    return targets


SPOOKY: int = 4


@dataclass
class Performance:
    win: int
    loss: int
    draw: int

    def rate(self) -> float:
        return float(self.win) / float(max(1, self.win + self.loss + self.draw))

    def __str__(self) -> str:
        global SPOOKY
        return f"+{str(self.win).zfill(SPOOKY)}-{str(self.loss).zfill(SPOOKY)}={str(self.draw).zfill(SPOOKY)}"


@dataclass(frozen=True)
class DataSplit:
    train_positions: Tensor
    train_targets: Tensor
    train_parity_signs: Tensor
    val_positions: Tensor
    val_targets: Tensor
    val_parity_signs: Tensor

    def batches(self, batch_size: int) -> Iterator[tuple[Tensor, Tensor, Tensor]]:
        order = torch.randperm(
            len(self.train_positions), device=self.train_positions.device
        )
        for start in range(0, len(order), batch_size):
            idx = order[start : start + batch_size]
            yield (
                self.train_positions[idx],
                self.train_targets[idx],
                self.train_parity_signs[idx],
            )


@dataclass(frozen=True)
class CandidateRecords:
    positions: Tensor  # [n, inp_dim] net input of the root state s_t
    leaf_positions: Tensor  # [n, inp_dim] net input of the PV leaf pv_t
    game_ids: Tensor  # [n]
    results: Tensor  # [n] game result R, black's perspective
    cand_results: Tensor  # [n] R from the recorded candidate's perspective
    parity_signs: Tensor  # [n] root parity
    leaf_parity_signs: Tensor  # [n] leaf parity
    leaf_terminals: Tensor  # [n] leaf is an exact terminal
    leaf_terminal_values: Tensor  # [n] exact terminal value, black's perspective

    def performance(self) -> Performance:
        result = Performance(0, 0, 0)
        for game_id in torch.unique(self.game_ids):
            in_game = self.game_ids == game_id
            game_result = self.cand_results[in_game][0].item()
            match game_result:
                case 1:
                    result.win += 1
                case -1:
                    result.loss += 1
                case _:
                    result.draw += 1
        return result

    def make_training_data(
        self, model: Module, training: CandidateTrainingConfig
    ) -> DataSplit:
        model.eval()
        with torch.no_grad():
            # Net(pv_t), black's perspective; exact terminals override the net
            leaf_values = model(self.leaf_positions) * self.leaf_parity_signs
            leaf_values = torch.where(
                self.leaf_terminals, self.leaf_terminal_values, leaf_values
            )

        targets = torch.empty_like(leaf_values)

        for game_id in torch.unique(self.game_ids):
            in_game = self.game_ids == game_id
            game_result = self.results[in_game][0].item()
            # self-pair games (i,i) record the candidate on every ply,
            # so smooth one chain per colour, not across players
            for sign in (1.0, -1.0):
                idxs = torch.nonzero(
                    in_game & (self.parity_signs == sign), as_tuple=True
                )[0]
                if len(idxs) == 0:
                    continue
                targets[idxs] = smooth(
                    leaf_values[idxs], game_result, training.smooth_factor
                )

        n_val = max(1, int(len(leaf_values) * training.val_frac))
        perm = torch.randperm(len(leaf_values), device=leaf_values.device)
        val_idx, train_idx = perm[:n_val], perm[n_val:]

        return DataSplit(
            train_positions=self.positions[train_idx],
            train_targets=targets[train_idx],
            train_parity_signs=self.parity_signs[train_idx],
            val_positions=self.positions[val_idx],
            val_targets=targets[val_idx],
            val_parity_signs=self.parity_signs[val_idx],
        )


@dataclass(frozen=True)
class PlayRecords:
    positions: list[Tensor]  # [candidate, n_i, inp_dim]
    leaf_positions: list[Tensor]  # [candidate, n_i, inp_dim]
    game_ids: list[Tensor]  # [candidate, n_i]
    results: list[Tensor]  # [candidate, n_i]
    cand_results: list[Tensor]  # [candidate, n_i]
    parity_signs: list[Tensor]  # [candidate, n_i]
    leaf_parity_signs: list[Tensor]  # [candidate, n_i]
    leaf_terminals: list[Tensor]  # [candidate, n_i]
    leaf_terminal_values: list[Tensor]  # [candidate, n_i]

    def index(self, index: int) -> CandidateRecords:
        return CandidateRecords(
            positions=self.positions[index],
            leaf_positions=self.leaf_positions[index],
            game_ids=self.game_ids[index],
            results=self.results[index],
            cand_results=self.cand_results[index],
            parity_signs=self.parity_signs[index],
            leaf_parity_signs=self.leaf_parity_signs[index],
            leaf_terminals=self.leaf_terminals[index],
            leaf_terminal_values=self.leaf_terminal_values[index],
        )

    @classmethod
    def load_merge(cls, files: list[Path], size: int) -> PlayRecords:
        # one JSON record per line:
        # {"result": R, "rows": [{"cand", "root": {...}, "leaf": {...}}]}
        inp_dim = NET_DIMS[size][0]
        n_cells = inp_dim - 2

        games: list[dict] = []
        for file in files:
            with open(file) as f:
                for line in f:
                    if line.strip():
                        games.append(json.loads(line))

        def board(row_part: dict) -> list[float]:
            return [*row_part["head"], *row_part["cells"]]

        n_cands = max(r["cand"] for g in games for r in g["rows"]) + 1
        fields = (
            "positions",
            "leaf_positions",
            "game_ids",
            "results",
            "cand_results",
            "parity_signs",
            "leaf_parity_signs",
            "leaf_terminals",
            "leaf_terminal_values",
        )
        per_cand = {c: {k: [] for k in fields} for c in range(n_cands)}
        for game_id, game in enumerate(games):
            result = float(game["result"])
            for cand in range(n_cands):
                rows = [r for r in game["rows"] if r["cand"] == cand]
                if not rows:
                    continue
                for part in ("root", "leaf"):
                    for r in rows:
                        if len(r[part]["cells"]) != n_cells:
                            raise ValueError(
                                f"{files[-1]}: {len(r[part]['cells'])} cells, "
                                f"expected {n_cells} for size {size}"
                            )
                parity = np.array([r["root"]["parity"] for r in rows], dtype=np.float32)
                per_cand[cand]["positions"].append(
                    torch.from_numpy(
                        np.array([board(r["root"]) for r in rows], dtype=np.float32)
                    )
                )
                per_cand[cand]["leaf_positions"].append(
                    torch.from_numpy(
                        np.array([board(r["leaf"]) for r in rows], dtype=np.float32)
                    )
                )
                per_cand[cand]["game_ids"].append(
                    torch.full((len(rows),), game_id, dtype=torch.int64)
                )
                per_cand[cand]["results"].append(
                    torch.full((len(rows),), result, dtype=torch.float32)
                )
                per_cand[cand]["cand_results"].append(torch.from_numpy(parity * result))
                per_cand[cand]["parity_signs"].append(torch.from_numpy(parity))
                per_cand[cand]["leaf_parity_signs"].append(
                    torch.from_numpy(
                        np.array([r["leaf"]["parity"] for r in rows], dtype=np.float32)
                    )
                )
                per_cand[cand]["leaf_terminals"].append(
                    torch.from_numpy(
                        np.array(
                            [r["leaf"]["terminal"] == 1 for r in rows],
                            dtype=np.bool_,
                        )
                    )
                )
                per_cand[cand]["leaf_terminal_values"].append(
                    torch.from_numpy(
                        np.array(
                            [r["leaf"]["terminal_value"] for r in rows],
                            dtype=np.float32,
                        )
                    )
                )

        def merged(cand: int, k: str, empty: Tensor) -> Tensor:
            blocks = per_cand[cand][k]
            return torch.concat(blocks, dim=0) if blocks else empty

        return cls(
            **{
                k: [
                    merged(
                        cand,
                        k,
                        torch.empty(
                            (0, inp_dim) if k.endswith("positions") else (0,),
                            dtype=torch.float32
                            if k not in ("game_ids", "leaf_terminals")
                            else (torch.int64 if k == "game_ids" else torch.bool),
                        ),
                    )
                    for cand in range(n_cands)
                ]
                for k in fields
            }
        )


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


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, x: Tensor) -> Tensor:
        # x: [n, 2 + n_cells]; the first 2 columns are the fixed head
        # inputs, the rest are raw cell states (0..6)
        n = x.shape[0]
        inp = torch.cat([x[:, :2], self.embedding(x[:, 2:].long()).view(n, -1)], dim=1)
        hidden = silu(self.hidden(inp))
        return tanh(self.output(hidden)).squeeze(-1)

    def fit(
        self, *, records: CandidateRecords, config: CandidateTrainingConfig
    ) -> tuple[float, float]:
        # training widgets
        optimizer = torch.optim.AdamW(self.parameters(), lr=config.lr)
        loss_fn = MSELoss()

        pbar = tqdm(range(config.epochs), desc="train", unit="ep")
        for _ in pbar:
            split = records.make_training_data(self, config)

            self.train()
            for positions, targets, parity_signs in split.batches(config.batch):
                optimizer.zero_grad()
                predictions = self(positions) * parity_signs
                loss = loss_fn(predictions, targets)
                loss.backward()
                optimizer.step()

            with torch.no_grad():
                self.eval()
                predictions = self(split.val_positions) * split.val_parity_signs
                mse = loss_fn(predictions, split.val_targets).item()
                corr = torch.corrcoef(torch.stack([predictions, split.val_targets]))[
                    0, 1
                ].item()

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

        return corr, mse

    def to_c(
        self,
        *,
        size: int,
        id: int | None,
        path: Path,
        note: str,
    ) -> None:
        upper = "FIVE" if size == 5 else "SIX"
        prefix = (f"cand_{id}_" if id is not None else "") + upper
        guard = f"WEIGHTS_{upper}" + (f"_{id}" if id 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",
                f"#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 setup(num_candidates: int, data_gen: DataGenConfig, root: Path) -> list[Net]:
    inp_dim, hidden = NET_DIMS[data_gen.size]
    candidates = [Net(inp_dim, hidden) for _ in range(num_candidates)]
    for id, cand in enumerate(candidates):
        cand.to_c(
            size=data_gen.size,
            id=id,
            path=root / "include" / f"cand_{id}_weights_{data_gen.size}.c",
            note=f"{id=}, v000, random init",
        )
        cand_copy_for(6 if data_gen.size == 5 else 5, id, root)
    return candidates


def run_step(
    *,
    candidates: list[Net],
    data_gen: DataGenConfig,
    iter: int,
    root: Path,
    training_configs: list[CandidateTrainingConfig],
):
    def data_path(i: int) -> Path:
        return root / "data" / f"selfplay-{data_gen.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=root, check=True, capture_output=True
        )

    # 2. generate self-play data
    data = data_path(iter)
    print(
        f"selfplay: {num_candidates=}, {data_gen.games=}, "
        f"{data_gen.size=}, {data_gen.depth=}, {data_gen.min_ply=}"
    )
    out_games: list[str] = []
    for out in fan_out(
        [
            [
                f"./selfplay-{i}-{j}",
                str(data_gen.size),
                str(data_gen.depth),
                str(data_gen.min_ply),
            ]
            for (i, j) in cand_pairs
        ],
        data_gen.games,
        data_gen.procs,
        "selfplay",
        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. compute performance
    data_files = [
        data_path(i) for i in range(max(iter - data_gen.replay_iters, 1), iter + 1)
    ]
    print("Loading data from", [str(c) for c in data_files], "...", end=" ")
    records = PlayRecords.load_merge(data_files, data_gen.size)
    print("candidates have", [f"{r.shape[0]:,}" for r in records.results], "rows.")
    results = [
        (i, cand, records.index(i).performance()) for i, cand in enumerate(candidates)
    ]

    print("Tournament results:")
    for idx, _, perf in sorted(results, key=lambda t: t[2].rate(), reverse=True):
        print(
            f"\tcandidate {idx:03d} {perf} {perf.rate() * 100.0:.2f}% {training_configs[idx]}"
        )

    # 4. write champ
    idx, champ, perf = max(results, key=lambda t: t[2].rate())
    path = root / "include" / f"weights_{data_gen.size}.c"
    champ.to_c(
        size=data_gen.size,
        id=None,
        path=path,
        note=f"{iter=}, perf={str(perf)}, rate={perf.rate() * 100.0:.3f}",
    )
    _ = subprocess.run(["make", "clean"], cwd=root, check=True, capture_output=True)
    _ = subprocess.run(["make", "native"], cwd=root, check=True, capture_output=True)
    print(f"Wrote candidate {idx:03d} to {str(path)}")

    # 5. fine-tune on the hybrid targets -> candidate
    for idx, cand in enumerate(candidates):
        corr, mse = cand.fit(records=records.index(idx), config=training_configs[idx])
        cand.to_c(
            size=data_gen.size,
            id=idx,
            path=root / "include" / f"cand_{idx}_weights_{data_gen.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",
    )

    args = p.parse_args()

    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,
    )
    tgsc = TrainingGridSearchConfig(
        smooth_factor=args.smooth_factor,
        epochs=args.epochs,
        batch=args.batch,
        lr=args.lr,
        val_frac=args.val_frac,
    )

    training_configs = tgsc.make_configs()

    (root / "data").mkdir(exist_ok=True)
    candidates = setup(
        num_candidates=len(training_configs), data_gen=data_gen, root=root
    )

    num_params = sum(p.numel() for p in candidates[0].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=data_gen,
            iter=iter,
            root=root,
            training_configs=training_configs,
        )
    print(f"\n== done, champion in include/weights_{args.size}.c ==")


if __name__ == "__main__":
    main()