"""Emit the MAE learning-rate sweep job spec, and measure the random-init floor.

Two patchings are swept, not one, because "MAE ranks mixtures worse than EAT"
and "16x16 patching handicaps MAE on a 64-mel input" are different claims and a
campaign launched before separating them cannot tell them apart later. Whichever
patching clears the floor by the wider margin is the primary arm; if they are
comparable, 16x16 wins on the grounds that it holds architecture fixed against
the EAT runs and leaves the objective as the sole variable.

The sweep runs at the FULL 1,912,024-clip horizon deliberately. The prior MAE
campaign's learning-rate.md established that a short sweep mis-ranks a long run,
and that mistake invalidated a whole batch there.
"""
import json
import subprocess
import sys
from pathlib import Path

sys.path.insert(0, "/workspace/code/eat-map-regmix")

CFG = "/workspace/code/eat-map-regmix/configs"
MAE = "/workspace/configs/mae.yaml"
DASHENG = "/workspace/configs/mae_dasheng_patch.yaml"
HORIZON = 1912024
LRS = [3e-4, 7e-4, 1.5e-3, 3e-3, 6e-3]

COMMON = [
    f"loop.budget_clips={HORIZON}",
    "loop.compile=false",          # compile=True hangs on this box; pinned for consistency
    "tracking.wandb_mode=disabled",
    "tracking.hub_upload=false",
    "tracking.campaign_id=mae-lr",
]

jobs = []
for tag, extra_cfg in [("p16", []), ("p4x64", [DASHENG])]:
    for lr in LRS:
        name = f"mae-lr/{tag}-lr{lr:g}"
        jobs.append({
            "name": name,
            "configs": [f"{CFG}/base.yaml", f"{CFG}/scale_15m.yaml",
                        f"{CFG}/budget_proxy.yaml", MAE] + extra_cfg,
            "sets": COMMON + [f"optim.lr={lr}", f"tracking.trial_id={tag}-lr{lr:g}"],
        })

out = Path("/workspace/scripts/jobs_lr.json")
out.write_text(json.dumps(jobs, indent=1))
print(f"wrote {out}: {len(jobs)} jobs "
      f"({len(LRS)} lrs x 2 patchings) at {HORIZON:,} clips")


# ---- random-init floor ------------------------------------------------------
# The gate every arm is measured against. It is a property of the architecture
# and the probe, not of the objective, so it must be measured once per patching.
def floor(tag: str, cfgs: list[str]) -> None:
    import torch
    from eatmap.config import load_config
    from eatmap.mae import MAEPretrainer
    from eatmap.runner import export_weights

    config = load_config(cfgs, [f"loop.budget_clips={HORIZON}", "loop.compile=false"])
    torch.manual_seed(0)
    model = MAEPretrainer(config)
    out_dir = Path(f"/workspace/runs/floor/{tag}/exports/step_00000000")
    export_weights(model, config, out_dir, step=0)
    subprocess.run(
        ["python", "-m", "eatmap.cli.probe", "--export", str(out_dir)],
        cwd="/workspace/code/eat-map-regmix", check=True,
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
    )
    result = json.loads((out_dir / "probe.json").read_text())
    print(f"  random-init floor [{tag}]: probe/map = {result['probe/map']:.5f}  "
          f"(holdout {result['probe/holdout_map']:.5f})")


print("\nmeasuring random-init floors (untrained encoder, same probe):")
floor("p16", [f"{CFG}/base.yaml", f"{CFG}/scale_15m.yaml", f"{CFG}/budget_proxy.yaml", MAE])
floor("p4x64", [f"{CFG}/base.yaml", f"{CFG}/scale_15m.yaml", f"{CFG}/budget_proxy.yaml", MAE, DASHENG])