"""
Rerun every experiment from scratch with the site's own exported code.

Same experiment grid as the original six sweeps. Every run trains with
bundle-template/btt (the code the playground exports) and is evaluated on:

    local-cpu        float32
    local-cpu-bf16   bfloat16 rounding emulated in PyTorch
    local-npu        the vendor toolkit. Its own report of what it ran on
                     (emulator or hardware) is saved with every result.

Results stream to runs.jsonl, one line per run, so the job can be stopped
and resumed. Usage:

    python experiments/rerun.py --workers 2
"""
import argparse
import json
import os
import platform
import sys
import time
from multiprocessing import Pool

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
sys.path.insert(0, os.path.join(ROOT, "bundle-template"))

BACKENDS = ["local-cpu", "local-cpu-bf16", "local-npu"]


def cfg(set_id, label, ls, gs, lr, ep, sched, noise=0.04, seed=0):
    return {
        "id": f"{set_id}:{label}",
        "set": set_id,
        "label": label,
        "config": {
            "schema": "btt.config/1", "target": "chirp", "layer_sizes": ls, "grid_sizes": gs,
            "lr": lr, "epochs": ep, "batch": 128, "noise_std": noise, "seed": seed,
            "scheduler": "cosine" if sched else "none", "backends": BACKENDS, "eval_points": 500,
        },
    }


def arch(w, d, g):
    ls = [1] + [w] * d + [1]
    return ls, [g] * (len(ls) - 1)


def jobs():
    out = []
    # Set A: width x depth x grid x learning rate, 8,000 steps, no schedule. 64 runs.
    for w in (6, 8, 10, 12):
        for d in (1, 2):
            for g in (8, 10, 12, 16):
                for lr in (0.003, 0.005):
                    out.append(cfg("A", f"w{w} d{d} g{g} lr{lr}", *arch(w, d, g), lr, 8000, False))
    # Set B: 30,000 steps, with and without a cosine schedule. 10 runs.
    for w, d, g, lr in ((8, 2, 8, 0.003), (12, 1, 12, 0.003), (12, 1, 16, 0.003), (8, 1, 12, 0.003), (8, 1, 8, 0.005)):
        for sched in (True, False):
            out.append(cfg("B", f"w{w} d{d} g{g} lr{lr} {'cosine' if sched else 'flat'}", *arch(w, d, g), lr, 30000, sched))
    # Set C: long runs, cosine schedule. 4 runs. The precision floor comes from here.
    for w, d, g, lr, ep, seed in ((12, 1, 12, 0.003, 60000, 0), (8, 2, 8, 0.003, 60000, 0),
                                  (16, 1, 16, 0.002, 60000, 0), (12, 1, 12, 0.003, 30000, 1)):
        out.append(cfg("C", f"w{w} d{d} g{g} lr{lr} {ep // 1000}k seed{seed}", *arch(w, d, g), lr, ep, True, seed=seed))
    # Set D: labelled architectures at lr 5e-3, 8,000 steps, no schedule. 8 runs.
    for label, ls, gs in (
        ("Starter", [1, 4, 1], [4, 4]),
        ("Wider, bigger grid", [1, 8, 1], [8, 8]),
        ("Wider still", [1, 16, 1], [16, 16]),
        ("Three layers", [1, 8, 8, 1], [8, 8, 8]),
        ("Three layers, wide", [1, 16, 16, 1], [16, 16, 16]),
        ("One hidden, huge grid", [1, 32, 1], [32, 32]),
        ("Big first grid", [1, 16, 16, 1], [32, 16, 16]),
        ("Four layers", [1, 16, 16, 16, 1], [16, 16, 16, 16]),
    ):
        out.append(cfg("D", label, ls, gs, 0.005, 8000, False))
    # Set E: the same family at lower learning rates, plus variants. 12 runs.
    for label, ls, gs, ep, lr, noise in (
        ("3-layer wide, lr 1e-3", [1, 16, 16, 1], [16, 16, 16], 8000, 0.001, 0.04),
        ("3-layer wide, lr 2e-3", [1, 16, 16, 1], [16, 16, 16], 8000, 0.002, 0.04),
        ("Huge grid, lr 1e-3", [1, 32, 1], [32, 32], 8000, 0.001, 0.04),
        ("Four layers, lr 1e-3", [1, 16, 16, 16, 1], [16, 16, 16, 16], 8000, 0.001, 0.04),
        ("3-layer 8, 20k steps", [1, 8, 8, 1], [8, 8, 8], 20000, 0.005, 0.04),
        ("2-layer 8, 20k steps", [1, 8, 1], [8, 8], 20000, 0.005, 0.04),
        ("2-layer 12, 12k steps", [1, 12, 1], [12, 12], 12000, 0.005, 0.04),
        ("3-layer, first grid 12", [1, 8, 8, 1], [12, 8, 8], 12000, 0.005, 0.04),
        ("3-layer 8, no noise", [1, 8, 8, 1], [8, 8, 8], 12000, 0.005, 0.0),
        ("3-layer 8, noise 0.10", [1, 8, 8, 1], [8, 8, 8], 12000, 0.005, 0.10),
        ("Narrow, fine grid", [1, 6, 6, 1], [10, 10, 10], 12000, 0.005, 0.04),
        ("Wide, coarse grid", [1, 12, 12, 1], [6, 6, 6], 12000, 0.005, 0.04),
    ):
        out.append(cfg("E", label, ls, gs, lr, ep, False, noise=noise))
    # Set F: longer runs with a cosine schedule and repeated seeds. 8 runs.
    for label, ls, gs, ep, lr, noise, seed in (
        ("2-layer 8, 30k", [1, 8, 1], [8, 8], 30000, 0.005, 0.04, 0),
        ("2-layer 8, 30k, seed 1", [1, 8, 1], [8, 8], 30000, 0.005, 0.04, 1),
        ("2-layer 8, 30k, seed 2", [1, 8, 1], [8, 8], 30000, 0.005, 0.04, 2),
        ("3-layer 8, 20k", [1, 8, 8, 1], [8, 8, 8], 20000, 0.003, 0.04, 0),
        ("3-layer 8, 20k, seed 1", [1, 8, 8, 1], [8, 8, 8], 20000, 0.003, 0.04, 1),
        ("3-layer 8, 20k, seed 2", [1, 8, 8, 1], [8, 8, 8], 20000, 0.003, 0.04, 2),
        ("3-layer 8, 20k, noise 0.02", [1, 8, 8, 1], [8, 8, 8], 20000, 0.003, 0.02, 0),
        ("3-layer 16, 12k, lr 1e-3", [1, 16, 16, 1], [16, 16, 16], 12000, 0.001, 0.04, 0),
    ):
        out.append(cfg("F", label, ls, gs, lr, ep, True, noise=noise, seed=seed))
    return out


def cost(j):
    c = j["config"]
    p = sum(c["layer_sizes"][i] * c["layer_sizes"][i + 1] * c["grid_sizes"][i] * 2 + c["layer_sizes"][i + 1]
            for i in range(len(c["layer_sizes"]) - 1))
    return c["epochs"] * (1.0 + p * 0.0008)


def work(j):
    os.environ["OMP_NUM_THREADS"] = "1"
    import torch
    torch.set_num_threads(1)
    from btt.runner import run
    t0 = time.time()
    res = run(j["config"], log_every=0)
    return {"id": j["id"], "set": j["set"], "label": j["label"], "wall_secs": round(time.time() - t0, 1),
            "host": {"machine": platform.machine(), "system": platform.system(), "node": "cloud-sandbox"},
            "result": res}


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--workers", type=int, default=2)
    ap.add_argument("--out", default=os.path.join(HERE, "runs.jsonl"))
    ap.add_argument("--only", help="comma separated set ids, e.g. A,C")
    a = ap.parse_args()
    done = set()
    if os.path.exists(a.out):
        with open(a.out) as fh:
            done = {json.loads(l)["id"] for l in fh if l.strip()}
    todo = [j for j in jobs() if j["id"] not in done and (not a.only or j["set"] in a.only.split(","))]
    todo.sort(key=cost, reverse=True)
    print(f"{len(jobs())} total, {len(done)} done, {len(todo)} to run", flush=True)
    with Pool(a.workers) as pool, open(a.out, "a") as fh:
        for n, r in enumerate(pool.imap_unordered(work, todo), 1):
            fh.write(json.dumps(r) + "\n"); fh.flush()
            b = r["result"]["backends"]
            print(f"[{n}/{len(todo)}] {r['id']:<44} {r['wall_secs']:>6}s  "
                  f"f32 {b['local-cpu']['mse']:.6f}  bf16 {b['local-cpu-bf16']['mse']:.6f}  "
                  f"tk {b['local-npu'].get('mse', float('nan')):.6f} ({b['local-npu'].get('device', {}).get('kind', '-')})",
                  flush=True)