#!/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 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(
    values: Tensor,
    result: float,
    td_lambda: float,
) -> Tensor:
    # values: V(s_t) for t = 0..T-1, shape [T]
    # result: final result R of the game

    # This is ema, computed backwards
    # target[T-1] = R
    # target[t]   = (1 - lambda) * V(s_{t+1}) + lambda * target[t+1]
    length = len(values)
    targets = torch.empty_like(values)
    targets[-1] = result
    for ply in range(length - 2, -1, -1):
        targets[ply] = (1.0 - td_lambda) * values[ply + 1] + td_lambda * targets[
            ply + 1
        ]
    return targets


@dataclass(frozen=True)
class SelfPlayRecords:
    positions: Tensor  #  [n, inp_dim]  feature vector of each position
    game_ids: Tensor  #  [n] id of the game the position belongs to
    plies: Tensor  # [n] ply number (position is before this move)
    results: Tensor  # [n] final result from black's perspective (+1 black won, -1 white won, 0 draw)

    @classmethod
    def load_selfplay(
        cls, fn: Path, size: int, device: torch.device
    ) -> SelfPlayRecords:
        # [ net input: 0-inp_dim ] [game_id] [ply] [game 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,
            ),
            plies=torch.tensor(
                df.iloc[:, inp_dim + 1].to_numpy(dtype=np.int64),
                device=device,
            ),
            results=torch.tensor(
                df.iloc[:, inp_dim + 2].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]),
            plies=torch.concat([spr.plies for spr in sprs]),
            results=torch.concat([spr.results 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_values = values[in_game]
            game_result = self.results[in_game][0].item()
            targets[in_game] = td_lambda_targets(game_values, game_result, td_lambda)

        return self.positions, targets


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


class Net(Module):
    def __init__(self, inp_dim: int, hidden: int) -> None:
        super().__init__()
        self.hidden = Linear(inp_dim, hidden, bias=True)
        self.output = Linear(hidden, 1, bias=True)

    def forward(self, x: Tensor) -> Tensor:
        hidden = torch.relu(self.hidden(x))
        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.Adam(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"

        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}_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
    cand_rate: float


def run_matches(
    *,
    size: int,
    depth: int,
    tournament_eps: float,
    games: int,
    procs: int,
    root: Path,
) -> TournamentResult:
    outs = list(
        fan_out(
            ["./match", str(size), f"{tournament_eps:.4f}", 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")
    aborts = len(winners) - champ - cand - draws
    played = len(winners) - aborts
    return TournamentResult(
        champ=champ,
        cand=cand,
        draws=draws,
        cand_rate=(cand + 0.5 * draws) / played if played else 0.0,
    )


# ===================================================================
# 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,
    eps_decay: float,
    eps_min: float,
    eps_start: float,
    games: int,
    iter: int,
    lr: float,
    min_ply: int,
    procs: int,
    promote_rate: float,
    replay_iters: int,
    root: Path,
    size: int,
    td_lambda: float,
    tournament_eps: float,
    tournament_games: int,
    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. self-play with the installed champion (include/weights_<size>.c)
    _ = subprocess.run(["make", "selfplay"], cwd=root, check=True, capture_output=True)

    # 2. generate self-play data (one game per process, in parallel)
    eps = max(eps_min, eps_start * (eps_decay ** (iter - 1)))
    csv = make_csv(iter)
    print(f"selfplay, eps={eps:.4f}, depth {depth}, {games} games on {procs} procs")
    out_rows: list[str] = []
    for out in fan_out(
        ["./selfplay", str(size), f"{eps:.4f}", 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")
    rows = len(out_rows)
    print(f"  {rows} rows")

    # 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", csvs)
    records = SelfPlayRecords.load_merge(csvs, size, device)
    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 linked)
    _ = subprocess.run(["make", "match"], cwd=root, check=True, capture_output=True)

    # 5. tournament + decision
    tournament_result = run_matches(
        root=root,
        size=size,
        depth=depth,
        tournament_eps=tournament_eps,
        games=tournament_games,
        procs=procs,
    )
    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, "
        + f"winrate={rate:.3f}, cand_rate={tournament_result.cand_rate:.3f}"
    )
    promote = decided >= int(0.25 * tournament_games) and rate >= promote_rate

    if promote:
        cand.to_c(
            size,
            root / "include" / f"weights_{size}.c",
            (f"v{iter:03d}, val_corr={corr:.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=10, help="max loop iterations")
    _ = p.add_argument(
        "--games", type=int, default=2000, 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(
        "--eps-start",
        type=float,
        default=0.15,
        help="initial self-play epsilon (random-move probability); kept close "
        "to --tour-eps so the value net trains on the policy it is deployed "
        "under",
    )
    _ = p.add_argument("--eps-min", type=float, default=0.05, help="epsilon floor")
    _ = p.add_argument(
        "--eps-decay", type=float, default=0.9, help="per-iteration epsilon decay"
    )
    _ = p.add_argument("--td-lambda", type=float, default=0.6, help="TD(lambda) lambda")
    _ = p.add_argument(
        "--epochs", type=int, default=20, help="candidate training epochs"
    )
    _ = p.add_argument("--lr", type=float, default=1e-3, 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(
        "--tournament-eps",
        type=float,
        default=0.05,
        help="tournament epsilon (random-move probability)",
    )
    _ = p.add_argument(
        "--tournament-games",
        type=int,
        default=None,
        help="tournament games (10/tour_eps if unset)",
    )
    _ = p.add_argument(
        "--promote",
        type=float,
        default=0.57,
        help="candidate win rate among non-drawn tournament games needed to promote (~p<0.03 one-sided binomial at the default 200 games)",
    )
    _ = 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 tournament",
    )
    _ = p.add_argument(
        "--min-ply", type=int, default=4, help="record positions at this ply or later"
    )

    args = p.parse_args()

    if args.tournament_games is None:
        args.tournament_games = max(20, int(round(10 / args.tournament_eps)))
    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)
    for iter in range(1, args.iters + 1):
        champ = run_step(
            batch=args.batch,
            champ=champ,
            depth=args.depth,
            epochs=args.epochs,
            eps_decay=args.eps_decay,
            eps_min=args.eps_min,
            eps_start=args.eps_start,
            games=args.games,
            iter=iter,
            lr=args.lr,
            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,
            tournament_eps=args.tournament_eps,
            tournament_games=args.tournament_games,
            val_frac=args.val_frac,
        )
    print(f"\n== done, champion in include/weights_{args.size}.c ==")


if __name__ == "__main__":
    main()