"""Entraîne un ou plusieurs bots à partir de zéro.

Exemples :
    python train.py evolution                  # un bot
    python train.py 3d peace --iters 300       # plusieurs bots
    python train.py all                        # les 6 bots qui apprennent, l'un après l'autre
    python train.py fly --resume               # reprendre depuis la dernière sauvegarde
    python train.py fly --resume --sans-limite # tourne jusqu'à ce que tu fasses Ctrl+C

Le meilleur cerveau (mesuré sur 300 cartes de contrôle jamais vues à l'entraînement)
est sauvegardé dans checkpoints/<bot>.npz ; la courbe d'apprentissage dans logs/<bot>.csv.
"""
from __future__ import annotations

import sys as _sys

try:  # console Windows : affiche les accents et les barres sans planter
    _sys.stdout.reconfigure(encoding="utf-8", errors="replace")
except Exception:
    pass

import argparse
import multiprocessing as mp
import os
import sys
import time

import numpy as np

from graphwar import bots
from graphwar.learn.common import SANS_FIN, Pool, validate
from graphwar.learn.evo import train_cem, train_es, train_ga
from graphwar.learn.ppo import train_ppo

TRAINERS = {"ga": train_ga, "es": train_es, "cem": train_cem, "ppo": train_ppo}
DEFAULT_ITERS = {"ga": 400, "es": 400, "cem": 400, "ppo": 300}


def saver(bot, which):
    path = os.path.join(bots.CKPT_DIR, f"{bot.key}{'' if which == 'best' else '_last'}.npz")
    bot.save(path)


def main():
    ap = argparse.ArgumentParser(description="Entraîne les bots Graphwar à partir de zéro.")
    ap.add_argument("bots", nargs="+", help=f"'all' ou parmi : {', '.join(bots.LEARNERS)}")
    ap.add_argument("--iters", type=int, default=None, help="nombre de générations / itérations")
    ap.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 2) - 2),
                    help="nombre de cœurs pour simuler les tirs (défaut : tous sauf 2)")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--resume", action="store_true", help="repartir de la dernière sauvegarde")
    ap.add_argument("--sans-limite", action="store_true",
                    help="ne s'arrête jamais tout seul : Ctrl+C pour arrêter (rien n'est perdu)")
    args = ap.parse_args()

    keys = bots.LEARNERS if "all" in args.bots else [k.lower() for k in args.bots]
    for k in keys:
        if k not in bots.LEARNERS:
            sys.exit(f"Bot inconnu ou qui n'apprend pas : {k}. Choix : {', '.join(bots.LEARNERS)}")

    pool = Pool(args.workers)
    print(f"Entraînement sur {pool.workers} cœur(s).\n")
    summary = []
    interrupted = False
    try:
        for k in keys:
            best0 = -1.0
            if args.resume:
                # on repart du meilleur cerveau sauvegardé (sinon du dernier état)
                best_path = os.path.join(bots.CKPT_DIR, f"{k}.npz")
                last = os.path.join(bots.CKPT_DIR, f"{k}_last.npz")
                bot = bots.make(k, path=best_path if os.path.exists(best_path) else last)
                if os.path.exists(os.path.join(bots.CKPT_DIR, f"{k}.npz")):
                    best0 = validate(pool, k, bots.make(k).policy.flat)["touche"]   # record à battre
            else:
                bot = bots.make(k, path="__aucun__", seed=args.seed)
                bot.trained_iters = 0
            iters = SANS_FIN if args.sans_limite else (args.iters or DEFAULT_ITERS[bot.learner])
            v0 = validate(pool, k, bot.policy.flat)
            print(f"━━ {bot.name} — {bot.method} — fonctions : {bot.vocab.name}")
            print(f"   cerveau : {' → '.join(map(str, bot.policy.sizes))}  ({bot.policy.n_params} poids)")
            print(f"   avant entraînement : {100 * v0['touche']:.1f}% de tirs qui touchent un ennemi\n")
            t0 = time.time()
            seed = args.seed + sum(map(ord, k)) + (int(time.time()) if args.resume else 0)   # nouvelles cartes en reprise
            if args.sans_limite:
                print("   Sans limite : Ctrl+C pour arrêter. Le record et l'état courant sont sauvegardés en continu.\n")
            best = TRAINERS[bot.learner](bot, pool, iters, np.random.default_rng(seed),
                                         save=saver, resume=args.resume, best0=best0)
            mins = (time.time() - t0) / 60
            print(f"\n   ✔ {bot.name} : record {100 * best:.1f}% en {mins:.1f} min → checkpoints/{k}.npz\n")
            summary.append((bot.name, v0["touche"], max(best, best0), mins))
    except KeyboardInterrupt:
        interrupted = True
        print("\n\nArrêté. Le meilleur cerveau est dans checkpoints/ ; relance avec --resume pour continuer.")
    finally:
        pool.close(force=interrupted)

    if summary:
        print("Bilan")
        for name, a, b, m in summary:
            print(f"  {name:11s} {100 * a:5.1f}%  →  {100 * b:5.1f}%   ({m:.1f} min)")
        print("\nPour les voir se battre :  python arena.py")


if __name__ == "__main__":
    mp.freeze_support()
    main()