"""Train the nnU-Net model for one pseudo-labeling round.

Epochs / initial lr / warm-start checkpoint are passed to NTUHLPLTrainer via
NNU_PL_* env vars (see scripts_nnu/trainers/ntuh_pl_trainer.py).

Usage:
  nnUNetv2_train wrapper: python scripts_nnu/03_nnu_train.py \
      --gpus 3 --epochs 250 --lr 1e-2 [--warmstart runs_nnu/round0/best_nnu.pth]
"""
import argparse
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from nnu_common import run, nnu_env, train_cmd, best_ckpt


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--gpus", type=int, default=3)
    ap.add_argument("--epochs", type=int, default=250)
    ap.add_argument("--lr", type=float, default=1e-2)
    ap.add_argument("--warmstart", default=None, help="checkpoint for full-weight warm start")
    ap.add_argument("--log", default=None)
    args = ap.parse_args()
    log = args.log or os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "logs", "nnu_train.log")
    run(train_cmd(args.gpus), log, env=nnu_env(epoch=args.epochs, lr=args.lr, warmstart=args.warmstart))
    print(f"[nnu:train] done; best checkpoint at {best_ckpt()}", flush=True)


if __name__ == "__main__":
    main()