#!/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
from math import log, ceil

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.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)


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


SPOOKY: int = 3


@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 CandidateRecords:
    positions: list[Tensor]  # [n, inp_dim]
    game_ids: list[Tensor]  # [n]
    results: list[Tensor]  # [n]
    cand_results: list[Tensor]  # [n]
    terminals: list[Tensor]  # [n]
    terminal_values: list[Tensor]  # [n]
    parity_signs: list[Tensor]  # [n]

    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, td_lambda: float
    ) -> tuple[Tensor, Tensor, Tensor]:
        model.eval()
        with torch.no_grad():
            values = model(self.positions) * self.parity_signs

        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, self.parity_signs


@dataclass(frozen=True)
class PlayRecords:
    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]
    terminals: list[Tensor]  # [candidate, n_i]
    terminal_values: list[Tensor]  # [candidate, n_i]
    parity_signs: list[Tensor]  # [candidate, n_i]

    def index(self, index: int) -> CandidateRecords:
        return CandidateRecords(
            positions=self.positions[index],
            game_ids=self.game_ids[index],
            results=self.results[index],
            cand_results=self.cand_results[index],
            terminals=self.terminals[index],
            terminal_values=self.terminal_values[index],
            parity_signs=self.parity_signs[index],
        )

    @classmethod
    def load_selfplay(cls, fn, size):
        df = pd.read_csv(fn, header=None)
        inp_dim = NET_DIMS[size][0]

        fields = {
            "parity_signs": df.iloc[:, 0].to_numpy(dtype=np.float32),
            "positions": df.iloc[:, 1 : inp_dim + 1].to_numpy(dtype=np.float32),
            "results": df.iloc[:, inp_dim + 1].to_numpy(dtype=np.float32),
            "terminals": df.iloc[:, inp_dim + 2].to_numpy(dtype=np.bool_),
            "terminal_values": df.iloc[:, inp_dim + 3].to_numpy(dtype=np.float32),
            "game_ids": df.iloc[:, inp_dim + 4].to_numpy(dtype=np.int64),
            "cand_results": df.iloc[:, inp_dim + 6].to_numpy(dtype=np.float32),
        }
        cand_ids = df.iloc[:, inp_dim + 5].to_numpy(dtype=np.int64)

        n_cands = int(cand_ids.max()) + 1
        gathered = {k: [] for k in fields}
        for cid in range(n_cands):
            mask = cand_ids == cid
            for k, arr in fields.items():
                gathered[k].append(torch.from_numpy(arr[mask]))

        return cls(**gathered)

    @classmethod
    def load_merge(cls, files: list[Path], size: int) -> PlayRecords:
        sprs = [PlayRecords.load_selfplay(file, size) for file in files]
        # for each field, gather [file0[i], file1[i], ...] for every index i
        # and concat along dim 0 -> merged list indexed by cand id
        return cls(
            positions=[
                torch.concat(vals, dim=0) for vals in zip(*(s.positions for s in sprs))
            ],
            game_ids=[
                torch.concat(vals, dim=0) for vals in zip(*(s.game_ids for s in sprs))
            ],
            results=[
                torch.concat(vals, dim=0) for vals in zip(*(s.results for s in sprs))
            ],
            cand_results=[
                torch.concat(vals, dim=0)
                for vals in zip(*(s.cand_results for s in sprs))
            ],
            terminals=[
                torch.concat(vals, dim=0) for vals in zip(*(s.terminals for s in sprs))
            ],
            terminal_values=[
                torch.concat(vals, dim=0)
                for vals in zip(*(s.terminal_values for s in sprs))
            ],
            parity_signs=[
                torch.concat(vals, dim=0)
                for vals in zip(*(s.parity_signs for s in sprs))
            ],
        )


# ===================================================================
# 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,
        td_lambda: float,
        val_frac: float,
        epochs: int,
        batch: int,
        lr: float,
    ) -> tuple[float, float]:
        # training widgets
        optimizer = torch.optim.AdamW(self.parameters(), lr=lr)
        loss_fn = MSELoss()

        pbar = tqdm(range(epochs), desc="train", unit="ep")
        for _ in pbar:
            # make data
            positions, targets, parity_signs = 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]
            val_parity_signs = parity_signs[val_idx]
            train_positions = positions[train_idx]
            train_targets = targets[train_idx]
            train_parity_signs = parity_signs[train_idx]

            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]) * train_parity_signs[batch_idx]
                )
                loss = loss_fn(predictions, train_targets[batch_idx])
                loss.backward()
                optimizer.step()

            with torch.no_grad():
                self.eval()
                predictions = self(val_positions) * val_parity_signs
                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}")

        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, size: int, root: Path) -> list[Net]:
    inp_dim, hidden = NET_DIMS[size]
    candidates = [Net(inp_dim, hidden) for _ in range(num_candidates)]
    for id, cand in enumerate(candidates):
        cand.to_c(
            size=size,
            id=id,
            path=root / "include" / f"cand_{id}_weights_{size}.c",
            note=f"{id=}, v000, random init",
        )
        cand_copy_for(6 if size == 5 else 5, id, root)
    return candidates


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

    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
    csv = csv_path(iter)
    print(f"selfplay: {num_candidates=}, {games=}, {size=}, {depth=}, {min_ply=}")
    out_rows: list[str] = []
    for out in fan_out(
        [
            [f"./selfplay-{i}-{j}", str(size), str(depth), str(min_ply)]
            for (i, j) in cand_pairs
        ],
        games,
        procs,
        "selfplay",
        root,
    ):
        out_rows.extend(out.splitlines())
    print(f"writing {len(out_rows):,} rows to {csv}")
    with open(csv, "w") as f:
        _ = f.write("\n".join(out_rows))
        if out_rows:
            _ = f.write("\n")

    # 3. compute performance
    csvs = [csv_path(i) for i in range(max(iter - replay_iters, 1), iter + 1)]
    print("Loading data from", [str(c) for c in csvs], "...", end=" ")
    records = PlayRecords.load_merge(csvs, 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}%")

    # 4. write champ
    idx, champ, perf = max(results, key=lambda t: t[2].rate())
    path = root / "include" / f"weights_{size}.c"
    champ.to_c(
        size=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 TD-leaf(lambda) targets -> candidate
    for idx, cand in enumerate(candidates):
        corr, mse = cand.fit(
            records=records.index(idx),
            td_lambda=td_lambda,
            val_frac=val_frac,
            epochs=epochs,
            batch=batch,
            lr=lr,
        )
        cand.to_c(
            size=size,
            id=idx,
            path=root / "include" / f"cand_{idx}_weights_{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="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=500,
        help="selfplay games per iteration",
    )
    _ = p.add_argument(
        "--num-candidates",
        type=int,
        default=3,
        help="number of distinct candidates to train",
    )
    _ = p.add_argument(
        "--replay-iters",
        type=int,
        default=0,
        help="number of previous data sources to retain",
    )
    _ = 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(
        "--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="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))

    global SPOOKY
    SPOOKY = ceil(
        log(
            max(1, args.games * args.num_candidates * (args.num_candidates + 1) // 2),
            10,
        )
    )

    (root / "data").mkdir(exist_ok=True)
    candidates = setup(num_candidates=args.num_candidates, size=args.size, 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(
            batch=args.batch,
            candidates=candidates,
            depth=args.depth,
            epochs=args.epochs,
            games=args.games,
            iter=iter,
            lr=args.lr,
            min_ply=args.min_ply,
            procs=args.procs,
            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()