# Concurrent paired training on the fork. Usage: python run.py PHASE STEPS EVAL_SPLIT ARMS_JSON
# ARMS_JSON: [{"arm":..., "opt":"adam|dimuon", "lr":..., "seed":..., "init":"<group>"}]; arms sharing "init" start from one saved state.
import sys, json, time, hashlib, numpy as np
from common import *
phase, steps, split, arms = sys.argv[1], int(sys.argv[2]), sys.argv[3], json.loads(sys.argv[4])
logf = open(f"{phase}_log.jsonl", "a")
def rec(**k):
    k["t"] = time.time(); logf.write(json.dumps(k, default=str) + "\n"); logf.flush(); print(json.dumps(k, default=str)[:400], flush=True)
T0 = time.time(); sc = service(); val, train, test = splits()
evald = [mk(r) for r in (val if split == "val" else test)]
rec(stage="eval_set", split=split, n=len(evald), sha256=batch_hash(evald))
mkc = lambda opt: retry(lambda: sc.create_lora_training_client(base_model=MODEL, rank=RANK, **opt_kw(opt)), "create")
def info(a, tc, restores):
    i = tc.get_info(); rec(stage="get_info", arm=a["arm"], restores=restores, optimizer=i.optimizer_config.type, config=i.optimizer_config.model_dump(), model=i.model_name, rank=i.lora_rank)
    want = "dimuon" if a["opt"] == "dimuon" else "adamw"
    assert i.optimizer_config.type == want and i.model_name == MODEL and i.lora_rank == RANK, "optimizer/model mismatch"
inits, C = {}, {}
for a in arms:
    tc = mkc(a["opt"]); info(a, tc, 0)
    if a["init"] not in inits:
        inits[a["init"]] = tc.save_state(f"{phase}-init-{a['init']}").result().path
        rec(stage="init_saved", group=a["init"], path=inits[a["init"]], by=a["arm"]); a["restores"] = 0
    else:
        tc.load_state(inits[a["init"]]).result(); a["restores"] = 1; info(a, tc, 1)
    o = tc.forward(evald[:4], loss_fn="cross_entropy").result()
    lp = np.concatenate([np.asarray(x["logprobs"].tolist(), dtype=np.float32) for x in o.loss_fn_outputs])
    rec(stage="initial_output", arm=a["arm"], init=a["init"], init_path=inits[a["init"]], restores=a["restores"], output_sha256=hashlib.sha256(lp.tobytes()).hexdigest())
    p0 = tc.save_state(f"{phase}-{a['arm']}-step0").result().path
    rec(stage="step0_saved", arm=a["arm"], path=p0, before_first_update=True)
    C[a["arm"]] = tc; a["order"] = train_order(train, a["seed"])
rec(stage="setup_done", seconds=time.time() - T0)
T1 = time.time()
for s in range(steps):
    ts = time.time(); futs = {}
    for a in arms:
        d = batch(a["order"], s); a["_d"] = d; a["_h"] = batch_hash(d)
        futs[a["arm"]] = (C[a["arm"]].forward_backward(d, loss_fn="cross_entropy"), None)
    for a in arms:
        f = futs[a["arm"]][0]; fb = retry(lambda: f.result(), "fb")
        p = step_params(a["opt"], a["lr"]); futs[a["arm"]] = (fb, C[a["arm"]].optim_step(p), type(p).__name__)
    for a in arms:
        fb, of, pt = futs[a["arm"]]; r = of.result(); ns, nt = token_nll(fb, a["_d"])
        rec(stage="train_step", arm=a["arm"], step=s + 1, batch_sha256=a["_h"], params_type=pt, train_nll_sum=ns, train_tokens=nt, train_token_mean_nll=ns / nt, grad_norm=r.metrics.get("grad_norm:mean"))
    print(f"step {s+1} {time.time()-ts:.1f}s", flush=True)
rec(stage="train_done", seconds=time.time() - T1, steps=steps, arms=len(arms))
T2 = time.time()
for a in arms:
    p = C[a["arm"]].save_state(f"{phase}-{a['arm']}-final").result().path
    s_, n_, per = eval_nll(C[a["arm"]], evald)
    rec(stage="eval", arm=a["arm"], split=split, step=steps, nll_sum=s_, n_tokens=n_, n_examples=len(per), token_mean=s_ / n_, example_mean=float(np.mean([x / max(y, 1) for x, y in per])), per_example=per, final_state=p)
rec(stage="eval_done", seconds=time.time() - T2, total_seconds=time.time() - T0)