#!/usr/bin/env python3
"""Resume-safe sequential launcher for the three 30-candidate HPO arms."""
import argparse, json, math, subprocess, sys, time, uuid
from pathlib import Path

ROOT=Path(__file__).resolve().parent.parent
def parse_args(argv=None):
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument("--output-dir",type=Path,required=True); p.add_argument("--resume",action="store_true")
    p.add_argument("--candidate-budget",type=int,default=30)
    p.add_argument("--n-epochs",type=int,default=1000); p.add_argument("--n-repeats",type=int,default=3)
    p.add_argument("--batch-size",type=int,default=32); p.add_argument("--num-workers",type=int,default=4)
    p.add_argument("--device",choices=("cpu","cuda"),default="cuda"); p.add_argument("--seed",type=int,default=42)
    p.add_argument("--no-wandb",action="store_true"); p.add_argument("--smoke",action="store_true")
    p.add_argument(
        "--validation-eval-policy",
        choices=("never", "on-source-improvement", "every-trial"),
        default="never",
    )
    p.add_argument("--plan",action="store_true",help="Print commands without creating state or running arms")
    return p.parse_args(argv)
def atomic(path,obj):
    t=path.with_suffix(".tmp"); t.write_text(json.dumps(obj,indent=2)+"\n"); t.replace(path)
def main(argv=None):
    a=parse_args(argv); a.output_dir.mkdir(parents=True,exist_ok=True); statep=a.output_dir/"orchestrator_state.json"
    if a.n_repeats != 3:
        raise ValueError("The three-arm comparison requires grouped CV=3 for every arm")
    if a.candidate_budget < 1:
        raise ValueError("candidate-budget must be positive")
    if statep.exists() and not a.resume: raise FileExistsError("pass --resume")
    state=json.loads(statep.read_text()) if statep.exists() else {"experiment_id":uuid.uuid4().hex[:8],"arms":{}}
    common=["--n-epochs",str(a.n_epochs),"--n-repeats",str(a.n_repeats),"--device",a.device,"--seed",str(a.seed)]
    population_size=6
    generations=math.ceil(a.candidate_budget/population_size)
    arms=[
      ("meta_evolution",[sys.executable,str(ROOT/"scripts/evolve_meta_model.py"),"--population-size",str(population_size),"--generations",str(generations),"--candidate-budget",str(a.candidate_budget),"--elite-count","2","--tournament-size","3","--no-hf-meta-policy-push"]),
      ("dataset_conditional",[sys.executable,str(ROOT/"scripts/run_optuna_comparison.py"),"--n-trials",str(a.candidate_budget),"--meta-hidden-size","64"]),
      ("global_shared",[sys.executable,str(ROOT/"scripts/run_global_shared_optuna.py"),"--n-trials",str(a.candidate_budget)])]
    if a.plan:
        for name,cmd in arms:
            print(name, " ".join(cmd))
        return 0
    for name,cmd in arms:
        rec=state["arms"].setdefault(name,{"status":"pending","wandb_run_name":f"three-arm-{state['experiment_id']}-{name}"})
        out=a.output_dir/name; cmd += ["--output-dir",str(out),"--wandb-run-name",rec["wandb_run_name"],*common]
        metadata_path=out/"run_metadata.json"
        completed_budget=0
        if metadata_path.exists():
            completed_budget=int(json.loads(metadata_path.read_text()).get("candidate_budget",0))
        if rec["status"]=="succeeded" and completed_budget >= a.candidate_budget:
            continue
        if name!="meta_evolution": cmd += ["--batch-size",str(a.batch_size),"--num-workers",str(a.num_workers)]
        else: cmd += ["--num-workers",str(a.num_workers)]
        if name == "dataset_conditional":
            cmd += ["--validation-eval-policy", a.validation_eval_policy]
        if a.no_wandb: cmd.append("--no-wandb")
        resume_markers = {
            "meta_evolution": (out / "checkpoint.npz", out / "state.json"),
            "dataset_conditional": (out / "run_metadata.json", out / "optuna.sqlite3"),
            "global_shared": (out / "run_metadata.json", out / "optuna.sqlite3"),
        }
        # An orchestrator retry is not necessarily an arm resume: preflight can
        # fail before the arm creates any checkpoint. Forward --resume only when
        # the arm has all of its own durable resume artifacts.
        if a.resume and all(path.exists() for path in resume_markers[name]):
            cmd.append("--resume")
        if a.smoke:
            cmd += ["--smoke-evaluator"] if name=="meta_evolution" else (["--smoke"] if name=="global_shared" else [])
            if name=="dataset_conditional": raise ValueError("dataset-conditional arm has no smoke evaluator; use its tests")
        rec.update(status="running",command=cmd,started_at=time.time()); atomic(statep,state)
        result=subprocess.run(cmd,cwd=ROOT)
        rec.update(status="succeeded" if result.returncode==0 else "failed",returncode=result.returncode,finished_at=time.time()); atomic(statep,state)
        if result.returncode: return result.returncode
    return 0
if __name__=="__main__": raise SystemExit(main())