"""
Run the revision experiments end to end (cross-platform; run from the S2F repository root).

    python -m experiments.run_all --legacy-root C:/Users/Kaveh/Documents/cellforce/dataset --steps audit
    python -m experiments.run_all --legacy-root ... --steps prepare,eval_released
    python -m experiments.run_all --steps train,eval_runs --variants s2f_full,unet,s2f_no_gan
    python -m experiments.run_all --legacy-root ... --stage1 C:/.../outputs/checkpoints/last/last_checkpoint_final.pth --steps eval_released

steps:
  audit          dataset audit only (no copying)                          -> experiments/results/dataset_audit/
  prepare        copy missing images into dataset/ + write experiments/splits/*.csv
  eval_released  released checkpoints (+ substrate sweep, mean-map baseline, optional stage-1)
                                                                          -> experiments/results/eval/
  train          train the listed variants (default: configs.PRIORITY), skipping finished runs.
                 One job per GPU in parallel (--gpus 0,1); the next variant starts as soon as a GPU is free.
                 Logs: experiments/runs/<variant>_seed<seed>/train.log
  eval_runs      evaluate every experiments/runs/*/best_checkpoint.pth
"""
import argparse
import json
import os
import subprocess
import sys
import time

from experiments.configs import PRIORITY, get_config
from experiments.layout import RUNS_DIR, S2F_ROOT


def sh(args):
    print("\n$ " + " ".join(args), flush=True)
    r = subprocess.run([sys.executable, "-m"] + args, cwd=S2F_ROOT)
    if r.returncode:
        sys.exit(r.returncode)


def finished(variant, seed):
    h = os.path.join(RUNS_DIR, f"{variant}_seed{seed}", "history.json")
    if not os.path.exists(h):
        return False
    return len(json.load(open(h))) >= get_config(variant)["epochs"]


def gpu_queue(queue, gpus, num_workers, extra=(), max_retries=2):
    """Run training jobs in parallel, one per GPU (each job sees only its GPU as cuda:0).
    A job that crashes (e.g. Windows access violation 3221225477) is re-queued and resumes from its last
    epoch checkpoint, up to `max_retries` times. Python is started with -X faulthandler so that a native
    crash writes the stack trace into train.log."""
    running, failed, tries = {}, [], {}
    queue = list(queue)
    while queue or running:
        for gpu in gpus:
            if gpu not in running and queue:
                v, seed = queue.pop(0)
                run_dir = os.path.join(RUNS_DIR, f"{v}_seed{seed}")
                os.makedirs(run_dir, exist_ok=True)
                log = open(os.path.join(run_dir, "train.log"), "a")
                env = dict(os.environ, CUDA_DEVICE_ORDER="PCI_BUS_ID", CUDA_VISIBLE_DEVICES=gpu, PYTHONUNBUFFERED="1")  # same numbering as nvidia-smi
                cmd = [sys.executable, "-X", "faulthandler", "-m", "experiments.train", "--variant", v, "--seed", seed,
                       "--resume", "--num-workers", num_workers] + list(extra)
                running[gpu] = (subprocess.Popen(cmd, cwd=S2F_ROOT, env=env, stdout=log, stderr=subprocess.STDOUT), v, seed, log)
                print(f"[{time.strftime('%H:%M')}] GPU {gpu}: started {v} seed {seed}  (log {run_dir}/train.log)", flush=True)
        time.sleep(10 if extra else 30)
        for gpu, (proc, v, seed, log) in list(running.items()):
            if proc.poll() is not None:
                log.close(); del running[gpu]
                if proc.returncode == 0:
                    print(f"[{time.strftime('%H:%M')}] GPU {gpu}: {v} seed {seed} finished", flush=True)
                    continue
                tries[(v, seed)] = tries.get((v, seed), 0) + 1
                if tries[(v, seed)] <= max_retries:
                    queue.insert(0, (v, seed))
                    print(f"[{time.strftime('%H:%M')}] GPU {gpu}: {v} seed {seed} crashed (exit {proc.returncode}); "
                          f"resuming from last checkpoint (retry {tries[(v, seed)]}/{max_retries})", flush=True)
                else:
                    failed.append((v, seed))
                    print(f"[{time.strftime('%H:%M')}] GPU {gpu}: {v} seed {seed} FAILED (exit {proc.returncode}, see train.log)", flush=True)
    if failed:
        print(f"failed jobs: {failed}")


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--legacy-root", default=None)
    ap.add_argument("--stage1", default=None, help="path to last_checkpoint_final.pth")
    ap.add_argument("--steps", default="audit")
    ap.add_argument("--variants", default=",".join(PRIORITY))
    ap.add_argument("--seeds", default="0")
    ap.add_argument("--gpus", default="0,1", help="GPU ids; one training job per GPU at a time")
    ap.add_argument("--num-workers", default="2", help="DataLoader workers per training job")
    ap.add_argument("--extra-train-args", default="", help="passed to experiments.train, e.g. '--epochs 1 --max-batches 2 --allow-cpu' for a smoke test")
    a = ap.parse_args()
    steps = a.steps.split(",")
    if ("audit" in steps or "prepare" in steps) and not a.legacy_root:
        sys.exit("--legacy-root is required for audit/prepare")
    if "audit" in steps:
        sh(["experiments.prepare_dataset", "--legacy-root", a.legacy_root])
    if "prepare" in steps:
        sh(["experiments.prepare_dataset", "--legacy-root", a.legacy_root, "--write"])
    if "eval_released" in steps:
        cmd = ["experiments.evaluate", "--released", "--sweep", "--meanmap"]
        if a.stage1:
            cmd += ["--stage1", a.stage1]
        sh(cmd)
    if "train" in steps:
        queue = [(v, s) for v in a.variants.split(",") for s in a.seeds.split(",") if not finished(v, int(s))]
        print(f"{len(queue)} training jobs on GPUs {a.gpus}: {queue}")
        gpu_queue(queue, a.gpus.split(","), a.num_workers, a.extra_train_args.split())
    if "eval_runs" in steps:
        sh(["experiments.evaluate", "--runs"])


if __name__ == "__main__":
    main()