"""Diagnose init-output mismatch: restore each rank's init twice in fresh sessions, fingerprint full fixed-batch logprobs."""
import json, time
from common import *
from research_exchange.restoration import output_fingerprint
d, b = fixed_datum()
P = {1: "tinker://858fb34c3b2a340a600975ef5aaa34c2/weights/run_b89b93886b-init-r1",
     32: "tinker://879b6d3bac21f78eff5ba092b3b3b86b/weights/run_b89b93886b-init-r32"}
out = {}
for R in (1, 32):
    for rep in range(2):
        sc = service_client(); tc = sc.create_training_client_from_state(P[R])
        lps = []
        for k in range(2):
            o = tc.forward([d], "cross_entropy").result()
            lps.append([float(x) for x in o.loss_fn_outputs[0]["logprobs"].tolist()])
        out[f"r{R}_session{rep}"] = {"model_id": tc.get_info().model_id, "fps": [output_fingerprint(l) for l in lps], "lp": lps[0]}
        print(R, rep, out[f"r{R}_session{rep}"]["model_id"], [f[:12] for f in out[f"r{R}_session{rep}"]["fps"]], flush=True)
        sc.close("success").result()
keys = list(out); base = out[keys[0]]["lp"]
for k in keys: print(k, "maxabs vs first", max(abs(a - c) for a, c in zip(base, out[k]["lp"])))
json.dump(out, open("diag_init.json", "w"))