from __future__ import annotations
import argparse, subprocess, copy, yaml, time, os, sys

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", required=True)
    ap.add_argument("--seeds", nargs="+", type=int, required=True)
    ap.add_argument("--steps", type=int, default=None)
    ap.add_argument("--eval_every", type=int, default=None)
    ap.add_argument("--tag_suffix", type=str, default="")
    args = ap.parse_args()

    base = yaml.safe_load(open(args.config, "r", encoding="utf-8"))
    for seed in args.seeds:
        cfg = copy.deepcopy(base)
        cfg["run"]["seed"] = int(seed)
        if args.tag_suffix:
            cfg["run"]["tag"] = f"{cfg['run']['tag']}_{args.tag_suffix}"
        if args.steps is not None:
            cfg["train"]["steps"] = int(args.steps)
        if args.eval_every is not None:
            cfg["train"]["eval_every"] = int(args.eval_every)

        tmp = f"/tmp/bihmoe_sweep_{int(time.time())}_seed{seed}.yaml"
        with open(tmp, "w", encoding="utf-8") as f:
            yaml.safe_dump(cfg, f, sort_keys=False)

        print(f"\n=== RUN seed={seed} cfg={tmp} ===")
        rc = subprocess.call(["uv", "run", "python", "scripts/train_poc_side_by_side.py", "--config", tmp])
        if rc != 0:
            print(f"ERROR: run failed with code {rc} for seed {seed}")
            sys.exit(rc)

if __name__ == "__main__":
    main()