#!/usr/bin/env -S uv run --script
#
# /// script
# requires-python = ">=3.14"
# dependencies = [
#   "pandas>=3.0.5",
#   "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

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

import numpy as np
import pandas as pd
import torch
from tqdm import tqdm

from torch.nn import Embedding, MSELoss, Linear, Module
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: (28, 64),  # FIVE_INP_NUM, FIVE_DENSE_NUM
    6: (39, 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)


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


def td_lambda_targets(
    leaf_values: Tensor,
    result: float,
    td_lambda: float,
) -> Tensor:
    targets = torch.empty_like(leaf_values)
    targets[-1] = result

    # G[t] = (1-lambda)V[t+1] + lambda G[t+1]
    for ply in range(len(leaf_values) - 2, -1, -1):
        targets[ply] = (1.0 - td_lambda) * leaf_values[ply + 1] + td_lambda * targets[
            ply + 1
        ]

    return targets


@dataclass(frozen=True)
class SelfPlayRecords:
    positions: Tensor  #  [n, inp_dim]
    game_ids: Tensor  #  [n] id of the game the position belongs to
    results: Tensor  # [n] final played result from black's perspective
    terminals: Tensor  # [n] true when the recorded position is a terminal leaf
    terminal_values: Tensor  # [n] exact value of a terminal leaf, otherwise 0

    @classmethod
    def load_selfplay(
        cls, fn: Path, size: int, device: torch.device
    ) -> SelfPlayRecords:
        # [net input] [game_id] [played result] [terminal flag] [leaf result]
        df = pd.read_csv(fn, header=None)
        inp_dim = NET_DIMS[size][0]
        return cls(
            positions=torch.tensor(
                df.iloc[:, 0:inp_dim].to_numpy(dtype=np.float32),
                dtype=torch.float32,
                device=device,
            ),
            game_ids=torch.tensor(
                df.iloc[:, inp_dim].to_numpy(dtype=np.int64),
                device=device,
            ),
            results=torch.tensor(
                df.iloc[:, inp_dim + 1].to_numpy(dtype=np.float32),
                dtype=torch.float32,
                device=device,
            ),
            terminals=torch.tensor(
                df.iloc[:, inp_dim + 2].to_numpy(dtype=np.bool_),
                dtype=torch.bool,
                device=device,
            ),
            terminal_values=torch.tensor(
                df.iloc[:, inp_dim + 3].to_numpy(dtype=np.float32),
                dtype=torch.float32,
                device=device,
            ),
        )

    @classmethod
    def load_merge(
        cls, files: list[Path], size: int, device: torch.device
    ) -> SelfPlayRecords:
        sprs = [SelfPlayRecords.load_selfplay(file, size, device) for file in files]
        return cls(
            positions=torch.concat([spr.positions for spr in sprs]),
            game_ids=torch.concat([spr.game_ids for spr in sprs]),
            results=torch.concat([spr.results for spr in sprs]),
            terminals=torch.concat([spr.terminals for spr in sprs]),
            terminal_values=torch.concat([spr.terminal_values for spr in sprs]),
        )

    def make_training_data(
        self, model: Module, td_lambda: float
    ) -> tuple[Tensor, Tensor]:
        model.eval()
        with torch.no_grad():
            values = model(self.positions)

        targets = torch.empty_like(values)

        for game_id in torch.unique(self.game_ids):
            in_game = self.game_ids == game_id
            game_terminals = self.terminals[in_game]
            game_terminal_values = self.terminal_values[in_game]
            game_values = values[in_game].clone()
            game_values[game_terminals] = game_terminal_values[game_terminals]
            game_result = self.results[in_game][0].item()
            game_targets = td_lambda_targets(game_values, game_result, td_lambda)
            # Terminal positions are also direct supervised examples. Do this
            # after constructing the TD sequence so it does not reset the EMA.
            game_targets[game_terminals] = game_terminal_values[game_terminals]
            targets[in_game] = game_targets

        return self.positions, targets


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

    def fit(
        self,
        *,
        self_play_records: SelfPlayRecords,
        td_lambda: float,
        val_frac: float,
        epochs: int,
        batch: int,
        lr: float,
    ) -> tuple[float, float]:
        # data
        positions, targets = self_play_records.make_training_data(self, td_lambda)
        n = len(positions)
        perm = torch.randperm(n, device=positions.device)
        n_val = max(1, int(n * val_frac))
        val_idx = perm[:n_val]
        train_idx = perm[n_val:]

        val_positions = positions[val_idx]
        val_targets = targets[val_idx]
        train_positions = positions[train_idx]
        train_targets = targets[train_idx]

        # training widgets
        optimizer = torch.optim.AdamW(self.parameters(), lr=lr)
        loss_fn = MSELoss()

        best_corr = -1.0
        best_state: dict[str, Tensor] | None = None
        best_mse = float("nan")

        pbar = tqdm(range(epochs), desc="train", unit="ep")
        for epoch in pbar:
            self.train()
            order = torch.randperm(len(train_positions), device=positions.device)
            for start in range(0, len(train_positions), batch):
                batch_idx = order[start : start + batch]
                optimizer.zero_grad()
                predictions = self(train_positions[batch_idx])
                loss = loss_fn(predictions, train_targets[batch_idx])
                loss.backward()
                optimizer.step()

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

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

            if corr > best_corr:
                best_corr = corr
                best_state = {
                    name: param.clone() for name, param in self.state_dict().items()
                }
                best_mse = mse

        if best_state is not None:
            self.load_state_dict(best_state)

        return best_corr, best_mse

    def to_c(
        self,
        size: int,
        path: Path,
        performance: tuple[str, ...] = (),
        cand: bool = False,
    ) -> None:
        prefix = ("cand_" if cand else "") + ("five" if size == 5 else "six")
        upper = "FIVE" if size == 5 else "SIX"

        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:
            _ = f.write("/* generated by nn_train.py, do not edit by hand */\n")
            for p in performance:
                _ = f.write(f"/* {p} */\n")
            _ = f.write(f'#include "weights_{size}.h"\n\n')
            _ = f.write(
                f"const float {prefix}_embedding[{len(E)}] =\n"
                f"{{\n{_fmt_flat(E, per_line=7)}\n}};\n\n"
            )
            _ = f.write(
                f"const float {prefix}_dense_weights[{upper}_DENSE_NUM]"
                f"[{upper}_INP_NUM] NN_ALIGNED =\n{{\n{_fmt_matrix(W)}\n}};\n\n"
            )
            _ = f.write(
                f"const float {prefix}_dense_biases[{upper}_DENSE_NUM] =\n"
                f"{{\n{_fmt_flat(B)}\n}};\n\n"
            )
            _ = f.write(
                f"const float {prefix}_logit_weights[{upper}_DENSE_NUM]"
                f" NN_ALIGNED =\n{{\n{_fmt_flat(LW)}\n}};\n\n"
            )
            _ = f.write(f"const float {prefix}_logit_bias = {LB:.9g};\n")


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


def fan_out(cmd: 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 fut in tqdm(
            concurrent.futures.as_completed(futs), total=n, desc=tag, unit="game"
        ):
            yield fut.result().stdout


def cand_copy_for(size: 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")
    prefix = "five" if size == 5 else "six"
    txt = src.read_text().replace(f"{prefix}_", f"cand_{prefix}_")
    _ = (root / "include" / f"cand_weights_{size}.c").write_text(txt)


# ===================================================================
# Tournament
# ===================================================================


@dataclass(frozen=True)
class TournamentResult:
    champ: int
    cand: int
    draws: int


def run_matches(
    *,
    depth: int,
    games: int,
    procs: int,
    root: Path,
    size: int,
    match_random_opening_length: int,
) -> TournamentResult:
    outs = list(
        fan_out(
            ["./match", str(size), f"{match_random_opening_length}", str(depth)],
            games,
            procs,
            "match",
            root,
        )
    )
    winners = [line.split()[0] for line in outs if line.strip()]
    champ = winners.count("current")
    cand = winners.count("candidate")
    draws = winners.count("draw")
    return TournamentResult(champ=champ, cand=cand, draws=draws)


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


def setup(size: int, root: Path) -> Net:
    inp_dim, hidden = NET_DIMS[size]
    champ = Net(inp_dim, hidden)
    champ.to_c(
        size,
        root / "include" / f"weights_{size}.c",
        ("v000, random init",),
    )
    cand_copy_for(6 if size == 5 else 5, root)
    return champ


def run_step(
    *,
    batch: int,
    champ: Net,
    depth: int,
    epochs: int,
    games: int,
    iter: int,
    lr: float,
    matches: int,
    match_random_opening_length: int,
    min_ply: int,
    procs: int,
    promote_rate: float,
    replay_iters: int,
    root: Path,
    size: int,
    td_lambda: float,
    val_frac: float,
) -> Net:
    def make_csv(i: int) -> Path:
        return root / "data" / f"selfplay-{size}-v{i:03d}.csv"

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

    # 1. rebuild
    _ = subprocess.run(["make", "selfplay"], cwd=root, check=True, capture_output=True)
    _ = subprocess.run(["make", "native"], cwd=root, check=True, capture_output=True)

    # 2. generate self-play data
    csv = make_csv(iter)
    print(f"selfplay: {size=}, {depth=}, {min_ply=}")
    out_rows: list[str] = []
    for out in fan_out(
        ["./selfplay", str(size), str(depth), str(min_ply)],
        games,
        procs,
        "selfplay",
        root,
    ):
        out_rows.extend(out.splitlines())
    with open(csv, "w") as f:
        _ = f.write("\n".join(out_rows))
        if out_rows:
            _ = f.write("\n")

    # 3. fine-tune the champion on TD(lambda) targets -> candidate
    device = next(champ.parameters()).device
    csvs = [make_csv(i) for i in range(max(iter - replay_iters, 1), iter + 1)]
    print("Loading data from", [str(c) for c in csvs], "...", end=" ")
    records = SelfPlayRecords.load_merge(csvs, size, device)
    print(f"{records.results.shape[0]:,} rows")
    cand = copy.deepcopy(champ)

    corr, mse = cand.fit(
        self_play_records=records,
        td_lambda=td_lambda,
        val_frac=val_frac,
        epochs=epochs,
        batch=batch,
        lr=lr,
    )
    cand.to_c(
        size,
        root / "include" / f"cand_weights_{size}.c",
        (f"v{iter:03d}, val_corr={corr:.4f}, val_mse={mse:.4f}",),
        cand=True,
    )

    # 4. build the match driver (champion + candidate)
    _ = subprocess.run(["make", "match"], cwd=root, check=True, capture_output=True)

    promote = True
    rate = -1.0
    if promote_rate > 0:
        # 5. tournament + decision
        print(f"match: {size=}, {match_random_opening_length=}, {depth=}")
        tournament_result = run_matches(
            depth=depth,
            games=matches,
            procs=procs,
            root=root,
            size=size,
            match_random_opening_length=match_random_opening_length,
        )
        decided = tournament_result.champ + tournament_result.cand
        rate = tournament_result.cand / decided if decided else 0.0
        print(
            f"champ {tournament_result.champ} - {tournament_result.cand} cand, {tournament_result.draws} draws, winrate={rate:.3f}"
        )
        promote = decided >= int(0.25 * matches) and rate > promote_rate

    if promote:
        cand.to_c(
            size,
            root / "include" / f"weights_{size}.c",
            (f"v{iter:03d}, val_corr={corr:.4f}, rate={rate:.4f}",),
        )
        print(f"== candidate v{iter:03d} PROMOTED (include/weights_{size}.c) ==")
        return cand
    print(f"== candidate v{iter:03d} rejected ==")
    return champ


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


def main() -> None:
    root = Path(__file__).resolve().parent.parent
    p = argparse.ArgumentParser(
        description="the TD(lambda) self-play bootstrap loop",
        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=1000,
        help="selfplay games per iteration",
    )
    _ = p.add_argument(
        "--replay-iters",
        type=int,
        default=4,
        help="number of previous self-play instances to mix in with current dataset during training",
    )
    _ = p.add_argument(
        "--td-lambda",
        type=float,
        default=0.7,
        help="TD(lambda) lambda",
    )
    _ = p.add_argument(
        "--epochs",
        type=int,
        default=20,
        help="candidate training epochs",
    )
    _ = p.add_argument(
        "--lr",
        type=float,
        default=1e-4,
        help="learning rate",
    )
    _ = p.add_argument(
        "--batch",
        type=int,
        default=128,
        help="training batch size",
    )
    _ = p.add_argument(
        "--val-frac",
        type=float,
        default=0.05,
        help="fraction of rows held out for validation",
    )
    _ = p.add_argument(
        "--matches",
        type=int,
        default=200,
        help="how many matches in a tournament",
    )
    _ = p.add_argument(
        "--match-random-opening-length",
        type=int,
        default=2,
        help="length of random opening in a match",
    )
    _ = p.add_argument(
        "--promote",
        type=float,
        default=0.53,
        help="candidate win rate among non-drawn matches needed to promote",
    )
    _ = p.add_argument(
        "--procs",
        type=int,
        default=None,
        help="parallel game processes (half the 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="record positions at this ply or later",
    )

    args = p.parse_args()

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

    (root / "data").mkdir(exist_ok=True)
    champ = setup(size=args.size, root=root)

    num_params = sum(p.numel() for p in champ.parameters() if p.requires_grad)
    print(f"Number of trainable parameters: {num_params:,}")

    for iter in range(1, args.iters + 1):
        champ = run_step(
            batch=args.batch,
            champ=champ,
            depth=args.depth,
            epochs=args.epochs,
            games=args.games,
            iter=iter,
            lr=args.lr,
            matches=args.matches,
            match_random_opening_length=args.match_random_opening_length,
            min_ply=args.min_ply,
            procs=args.procs,
            promote_rate=args.promote,
            replay_iters=args.replay_iters,
            root=root,
            size=args.size,
            td_lambda=args.td_lambda,
            val_frac=args.val_frac,
        )
    print(f"\n== done, champion in include/weights_{args.size}.c ==")


if __name__ == "__main__":
    main()