# Correctness + restoration preflight on the fork; all clients start from one saved init state P0.
import json, time, hashlib, numpy as np, tinker
from common import *
log = []
def rec(**k):
    k["t"] = time.time(); log.append(k); print(json.dumps(k, default=str)[:600], flush=True)
    json.dump(log, open("preflight_log.json", "w"), indent=1, default=str)
def info(tc, arm, after_restore):
    i = tc.get_info()
    rec(stage="get_info", arm=arm, after_restore=after_restore, optimizer=i.optimizer_config.type if i.optimizer_config else None,
        actual_config=i.optimizer_config.model_dump() if i.optimizer_config else None, model=i.model_name, rank=i.lora_rank, model_id=i.model_id)
def out_hash(tc, data, arm):
    o = tc.forward(data, loss_fn="cross_entropy").result()
    lps = [np.asarray(x["logprobs"].tolist(), dtype=np.float32) for x in o.loss_fn_outputs]
    np.save(f"pf_initial_logprobs_{arm}.npy", np.concatenate(lps))
    rec(stage="initial_output", arm=arm, output_sha256=hashlib.sha256(np.concatenate(lps).tobytes()).hexdigest(), batch_sha256=batch_hash(data))
sc = service(); val, train, test = splits(); order = train_order(train, 0)
b1, b2 = batch(order, 0), batch(order, 1); probe = b1[:4]
rec(stage="fixed_batch", batch1_sha256=batch_hash(b1), batch2_sha256=batch_hash(b2))
mk_client = lambda opt: retry(lambda: sc.create_lora_training_client(base_model=MODEL, rank=RANK, **opt_kw(opt)), "create")
I = mk_client("dimuon"); info(I, "dimuon_lr1e-3_uninterrupted", 0)
P0 = I.save_state("pf-init").result().path; rec(stage="save", name="pf-init", path=P0)
out_hash(I, probe, "dimuon_I")
def fresh(opt, arm):
    tc = mk_client(opt); info(tc, arm, 0); tc.load_state(P0).result(); info(tc, arm, 1); out_hash(tc, probe, arm); return tc
def step(tc, data, params, name, arm):
    fb = tc.forward_backward(data, loss_fn="cross_entropy").result()
    r = tc.optim_step(params).result()
    p = tc.save_state(name).result().path
    rec(stage="step", arm=arm, name=name, path=p, params=params.model_dump(), params_type=type(params).__name__, fb=fb.metrics, grad_norm=r.metrics.get("grad_norm:mean"))
G = fresh("adam", "grad_recovery_adam")
step(G, b1, tinker.AdamParams(learning_rate=1e6, beta1=0.9, beta2=0.95, eps=1e6), "pf-grad", "grad_recovery_adam")
A = fresh("adam", "adamw_lr1e-4")
step(A, b1, step_params("adam", 1e-4), "pf-adam-1e-4", "adamw_lr1e-4")
for lr in (1e-4, 1e-2):
    D = fresh("dimuon", f"dimuon_lr{lr:g}")
    step(D, b1, step_params("dimuon", lr), f"pf-dim-{lr:g}", f"dimuon_lr{lr:g}")
step(I, b1, step_params("dimuon", 1e-3), "pf-dim-0.001", "dimuon_lr1e-3_uninterrupted")
step(I, b2, step_params("dimuon", 1e-3), "pf-U2", "dimuon_lr1e-3_uninterrupted")
R = mk_client("dimuon"); info(R, "dimuon_restored", 0)
R.load_state_with_optimizer([e for e in log if e.get("name") == "pf-dim-0.001"][0]["path"]).result(); info(R, "dimuon_restored", 1)
step(R, b2, step_params("dimuon", 1e-3), "pf-R2", "dimuon_restored")
rec(stage="done")