"""Build the canonical distributions file and a 4-job pre-campaign validation wave.

Two questions, one 20-minute wave on 4 GPUs, before 5 hours are committed:

1. Is lr 7e-4 actually the peak? The sweep bracketed it at 3e-4 / 7e-4 / 1.5e-3
   (0.1665 / 0.1826 / 0.1527), so 5e-4 and 1e-3 refine the interior.
2. Does that lr survive a concentrated mixture? It was tuned on natural, but the
   64 arms span 4.39-7.47x repetition and 284k-808k distinct clips. If the
   extremes collapse at 7e-4, a slice of the campaign is noise and the ranking is
   contaminated in a way no amount of analysis recovers. d7 has the fewest live
   clusters (6 of 20) and d185 the fewest distinct clips (293,514) -- the two
   corners of the design.

The distributions file uses the record form ({dist_id, weights}) because
data.build_mixture_plan matches dist_id on the field rather than list position,
and these 64 ids are a sparse subset of 0..255.
"""
import json
from pathlib import Path

CFG = "/workspace/code/eat-map-regmix/configs"
MAE = "/workspace/configs/mae.yaml"
HORIZON = 1912024
CLUSTER_INDEX = "/workspace/data/asmel_flat/train.cluster_index.npy"
DIST = "/workspace/artifacts/canonical64_distributions.json"
BEST_LR = 7e-4

spec = json.loads(Path("/workspace/analysis/canonical64.json").read_text())

# ---- canonical distributions file ------------------------------------------
records = [{"dist_id": v["dist_id"], "weights": v["weights"], "batch": v["batch"]}
           for v in sorted(spec.values(), key=lambda r: r["dist_id"])]
Path(DIST).write_text(json.dumps({
    "source": "regmix32-base-1ep + regmix32b-base-1ep realized weights",
    "note": "the 64-mixture canonical pool; cap disabled so the realized vector "
            "is identical at every budget and objective",
    "distributions": records,
}, indent=1))
print(f"wrote {DIST}: {len(records)} records, "
      f"dist_id range {records[0]['dist_id']}-{records[-1]['dist_id']}")

COMMON = [
    f"loop.budget_clips={HORIZON}",
    "loop.compile=false",
    "tracking.wandb_mode=disabled",
    "tracking.hub_upload=false",
    "tracking.campaign_id=mae-validate",
]
MIX = [
    f"mixture.cluster_index={CLUSTER_INDEX}",
    f"mixture.distributions={DIST}",
    # Cap off: at this budget these vectors need 4.39-7.47x repetition, and
    # capping would change the realized vector per budget, so different cells
    # would be ranking different mixtures.
    "mixture.repetition_cap=1000.0",
]

jobs = []
for lr in (5e-4, 1e-3):  # refine the interior optimum
    jobs.append({
        "name": f"mae-validate/natural-lr{lr:g}",
        "configs": [f"{CFG}/base.yaml", f"{CFG}/scale_15m.yaml", f"{CFG}/budget_proxy.yaml", MAE],
        "sets": COMMON + [f"optim.lr={lr}", f"tracking.trial_id=natural-lr{lr:g}"],
    })
for did in (7, 185):  # the two corners of the concentration design
    jobs.append({
        "name": f"mae-validate/d{did}-lr{BEST_LR:g}",
        "configs": [f"{CFG}/base.yaml", f"{CFG}/scale_15m.yaml", f"{CFG}/budget_proxy.yaml", MAE],
        "sets": COMMON + MIX + [f"optim.lr={BEST_LR}", f"mixture.dist_id={did}",
                                f"tracking.trial_id=d{did}-lr{BEST_LR:g}"],
    })

out = Path("/workspace/scripts/jobs_validate.json")
out.write_text(json.dumps(jobs, indent=1))
print(f"wrote {out}: {len(jobs)} jobs")
for j in jobs:
    print(f"  {j['name']}")