#
# 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 os
import subprocess
from pathlib import Path

import numpy as np

from ctak_train.candidate import Candidate, DataGenConfig, TrainingGridSearchConfig
from ctak_train.data import GameLog, load_window
from ctak_train.net import NET_DIMS
from ctak_train.selfplay import fan_out
from ctak_train.strength import fit_strengths, plot_strengths


def run_step(
    *,
    candidates: list[Candidate],
    data_gen_config: DataGenConfig,
    log: GameLog,
    strength_history: list[np.ndarray],
    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. load the training window and append this iteration's games to the log
    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, new_games = load_window(
        data_files, data_gen_config.size, num_candidates, iter
    )
    print("candidates have", [f"{p.n_rows:,}" for p in pools], "rows.")
    log.append(new_games)

    # 4. fit strengths over the whole history, and plot the trajectory
    strengths, se = fit_strengths(log, num_candidates)
    strength_history.append(strengths)
    plot_strengths(
        strength_history,
        candidates,
        data_gen_config.root / "data" / "candidate_strength.png",
    )

    rated_games = int(np.count_nonzero(log.black != log.white))
    ranked = sorted(range(num_candidates), key=lambda i: strengths[i], reverse=True)
    print(
        f"Strength fit ({rated_games:,} head-to-head games, "
        f"iterations 1..{iter}, {rated_games // num_candidates:,} games/candidate):"
    )
    for idx in ranked:
        print(
            f"\tcandidate {idx:03d}: strength {strengths[idx]:6.1f} "
            f"± {1.96 * se[idx]:5.1f} Config {candidates[idx]}"
        )

    # 5. write champ, unconditionally
    idx = ranked[0]
    path = data_gen_config.root / "include" / f"weights_{data_gen_config.size}.c"
    candidates[idx].net.to_c(
        size=data_gen_config.size,
        ident=None,
        path=path,
        note=f"{iter=}, strength={strengths[idx]:.1f}",
    )
    _ = 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}")

    # 6. 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}",
        )


def main() -> None:
    # repo root, relative to this file at <root>/train/src/ctak_train/nn_train.py
    root = Path(__file__).resolve().parents[3]
    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="number of previous data files to keep in the training pools "
        "(every game record is kept for the strength fit)",
    )
    _ = 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)",
    )

    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,
    )
    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:,}")

    log = GameLog.empty()
    strength_history: list[np.ndarray] = []

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


if __name__ == "__main__":
    main()