"""Run the official 10-fold training matrix.

Every ``(configuration, fold)`` pair is trained with the same pinned schedule.
A run counts as complete only when it wrote ``summary.json`` - Ultralytics
emits ``weights/best.pt`` from the first epoch, so a checkpoint alone does not
mean the schedule finished. The matrix is therefore resumable across
invocations without silently keeping a half-trained fold.

Usage
-----
    python scripts/03_train.py                      # everything still missing
    python scripts/03_train.py --groups main        # one group
    python scripts/03_train.py --exp Exp4_2p5D_Strict --folds 0
    python scripts/03_train.py --dry-run            # just print the plan
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from luna_rev import config as cfg
from luna_rev import splits
from luna_rev.train import is_complete

LOG_PATH = cfg.RESULTS_DIR / "training_log.jsonl"

#: Seconds to idle after each run so Windows can reap the dataloader workers
#: before the next run starts allocating.
COOLDOWN_S = 8


def run_in_subprocess(exp, fold, force: bool = False, epochs: int | None = None) -> dict:
    """Train one pair in a fresh interpreter and return its record.

    Isolating each run guarantees that workers, pinned buffers and the CUDA
    context are released before the next one begins; sharing one interpreter
    across 120 runs exhausted page-locked host memory partway through.
    """
    import subprocess

    cmd = [sys.executable, "-m", "luna_rev._train_worker", exp.name, str(fold.index)]
    if force:
        cmd.append("--force")
    if epochs:
        cmd += ["--epochs", str(epochs)]

    proc = subprocess.run(cmd, cwd=str(cfg.ROOT), capture_output=True, text=True,
                          encoding="utf-8", errors="replace")
    time.sleep(COOLDOWN_S)

    for line in reversed((proc.stdout or "").splitlines()):
        if line.startswith("TRAIN_RECORD "):
            return json.loads(line[len("TRAIN_RECORD "):])

    tail = "\n".join((proc.stderr or "").strip().splitlines()[-12:])
    print(f"    subprocess exit {proc.returncode}\n{tail}", flush=True)
    return {"exp": exp.name, "fold": fold.index, "status": "failed",
            "returncode": proc.returncode, "stderr_tail": tail}


def build_plan(groups, exp_names, fold_filter) -> list[tuple]:
    """Expand the requested selection into ``(exp, fold)`` pairs."""
    folds = {f.index: f for f in splits.get_folds("official")}
    experiments = cfg.experiments_for_group(*groups) if groups else list(cfg.ALL_EXPERIMENTS)
    if exp_names:
        wanted = set(exp_names)
        experiments = [e for e in experiments if e.name in wanted]
        unknown = wanted - {e.name for e in experiments}
        if unknown:
            raise SystemExit(f"Unknown experiment(s): {sorted(unknown)}")

    plan = []
    for exp in experiments:
        for k in cfg.folds_for(exp):
            if fold_filter is not None and k not in fold_filter:
                continue
            plan.append((exp, folds[k]))
    return plan


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--groups", nargs="*", default=None,
                    help="main r_sweep w_sweep neg_ablation seeds yolo26")
    ap.add_argument("--exp", nargs="*", default=None, help="explicit configuration names")
    ap.add_argument("--folds", nargs="*", type=int, default=None)
    ap.add_argument("--force", action="store_true", help="retrain even if weights exist")
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--epochs", type=int, default=None, help="override the schedule (benchmarking)")
    args = ap.parse_args()

    plan = build_plan(args.groups, args.exp, set(args.folds) if args.folds else None)
    pending = [(e, f) for e, f in plan if args.force or not is_complete(e, f.index)]

    print(f"planned runs : {len(plan)}")
    print(f"already done : {len(plan) - len(pending)}")
    print(f"to train     : {len(pending)}")
    if args.dry_run:
        for e, f in pending:
            print(f"  {e.name:28s} fold{f.index}  [{e.model}, r={e.r_sample}, "
                  f"w_min={e.w_min_px}, neg={e.negatives}, seed={e.seed}]")
        return 0
    if not pending:
        return 0

    t_start = time.time()
    failures: list[tuple] = []
    for i, (exp, fold) in enumerate(pending, 1):
        t0 = time.time()
        print(f"\n[{i}/{len(pending)}] {exp.name} fold{fold.index} "
              f"(elapsed {(time.time() - t_start) / 3600:.2f} h)", flush=True)
        rec = run_in_subprocess(exp, fold, force=args.force, epochs=args.epochs)
        rec["wall_seconds"] = round(time.time() - t0, 1)
        with LOG_PATH.open("a", encoding="utf-8") as fh:
            fh.write(json.dumps(rec) + "\n")
        if rec["status"] == "failed":
            failures.append((exp, fold))
            print(f"    -> FAILED after {rec['wall_seconds'] / 60:.1f} min "
                  f"(continuing; will retry at the end)", flush=True)
            continue
        rate = (time.time() - t_start) / i
        print(f"    -> {rec['status']} in {rec['wall_seconds'] / 60:.1f} min "
              f"| ETA {(len(pending) - i) * rate / 3600:.1f} h", flush=True)

    # One retry pass: the failures we have seen are transient resource
    # exhaustion, which a fresh process after an idle gap usually clears.
    if failures:
        print(f"\nRetrying {len(failures)} failed run(s) ...", flush=True)
        still_failing = []
        for exp, fold in failures:
            time.sleep(30)
            rec = run_in_subprocess(exp, fold, force=True, epochs=args.epochs)
            with LOG_PATH.open("a", encoding="utf-8") as fh:
                fh.write(json.dumps({**rec, "retry": True}) + "\n")
            if rec["status"] == "failed":
                still_failing.append(f"{exp.name} fold{fold.index}")
        if still_failing:
            print("PERMANENT FAILURES: " + ", ".join(still_failing), flush=True)
            return 1

    print(f"\nAll done in {(time.time() - t_start) / 3600:.2f} h")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())