"""Dry-run-by-default launcher for BASELINE_SPEC.md's first BP twin rung."""
from __future__ import annotations

import argparse
import subprocess
import sys


VOCAB = 32768
LAYERS = 12
WIDTH = 768
HEADS = 12
CONTEXT = 256
BATCH = 24
PARAMETERS = 135303936
FULL_STEPS_ARG = 440000  # Match fw135m_bsign; casc_bp_train.py loops inclusively: 440,001 updates.
LR_SWEEP = ("7e-4", "1e-3", "1.4e-3")


def parameter_count(vocab=VOCAB, layers=LAYERS, width=WIDTH):
    hidden = ((8 * width // 3) + 63) // 64 * 64
    return 2 * vocab * width + width + layers * (
        4 * width * width + 3 * width * hidden + 4 * width
    )


def command(tag, lr, steps, warmup, seed, wandb_project):
    return [
        sys.executable,
        "casc_bp_train.py",
        "--tag",
        tag,
        "--L",
        str(LAYERS),
        "--C",
        str(WIDTH),
        "--H",
        str(HEADS),
        "--T",
        str(CONTEXT),
        "--B",
        str(BATCH),
        "--steps",
        str(steps),
        "--lr",
        lr,
        "--warmup",
        str(warmup),
        "--amp",
        "--olmo2",
        "--wd",
        "0.1",
        "--opt",
        "muon",
        "--muon_lr",
        "0.02",
        "--cosine",
        "--lr_min_ratio",
        "0.1",
        "--data",
        "fineweb_edu",
        "--seed",
        str(seed),
        "--save_every",
        "5000",
        "--log",
        "100",
        "--wandb",
        wandb_project,
        "--wandb_run",
        tag if wandb_project else "",
    ]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--mode", choices=("smoke", "sweep", "full"), default="smoke")
    ap.add_argument("--execute", action="store_true", help="run commands; default only prints")
    ap.add_argument("--lr", default="1e-3", help="selected LR for --mode full")
    ap.add_argument("--seed", type=int, default=1)
    ap.add_argument("--sweep_steps", type=int, default=FULL_STEPS_ARG)
    ap.add_argument("--wandb_project", default="")
    args = ap.parse_args()

    count = parameter_count()
    if count != PARAMETERS:
        raise RuntimeError(f"parameter formula returned {count}, expected {PARAMETERS}")

    if args.mode == "smoke":
        runs = [(f"fw135m_bp_smoke_s{args.seed}", args.lr, 400, 50)]
    elif args.mode == "sweep":
        runs = [
            (
                f"fw135m_bp_lr{lr.replace('-', 'm').replace('.', 'p')}_s{args.seed}",
                lr,
                args.sweep_steps,
                1000,
            )
            for lr in LR_SWEEP
        ]
    else:
        runs = [(f"fw135m_bp_s{args.seed}", args.lr, FULL_STEPS_ARG, 1000)]

    print(
        f"# L{LAYERS} C{WIDTH} H{HEADS} T{CONTEXT} B{BATCH} "
        f"| {PARAMETERS:,} params | target 20N={20 * PARAMETERS:,} tokens",
        flush=True,
    )
    for tag, lr, steps, warmup in runs:
        cmd = command(tag, lr, steps, warmup, args.seed, args.wandb_project)
        print(" ".join(f'"{item}"' if item == "" else item for item in cmd), flush=True)
        if args.execute:
            subprocess.run(cmd, check=True)


if __name__ == "__main__":
    main()