"""Optimizer preflight + server baseline parity on a fixed real-rollout batch."""
import json, time, hashlib, concurrent.futures as cf
import numpy as np, torch
from tinker import types
from common import service_client
import gsm8k_eval as ge, rl_score_centering as rl

out = {}
sc = service_client(); pool = cf.ThreadPoolExecutor(64)
t0 = time.time()
tc = rl.create_client(sc); out["t_create"] = time.time() - t0
info = tc.get_info(); out["client_info"] = json.loads(info.model_dump_json()) if hasattr(info, "model_dump_json") else str(info)
# Optimizer preflight: one example, cross-entropy, one AdamParams(lr=1e-4) step.
toks = ge.prompt_tokens("What is 2+3?") + ge.tokenizer().encode("2+3=5\n#### 5")
d = types.Datum(model_input=types.ModelInput.from_ints(toks[:-1]), loss_fn_inputs={
    "target_tokens": types.TensorData.from_numpy(np.array(toks[1:], dtype=np.int64)),
    "weights": types.TensorData.from_numpy(np.ones(len(toks) - 1, dtype=np.float32))})
fb = tc.forward_backward([d], "cross_entropy").result()
opt = tc.optim_step(types.AdamParams(learning_rate=rl.LR)).result()
out["optimizer_preflight"] = {"adam_params": json.loads(types.AdamParams(learning_rate=rl.LR).model_dump_json()),
                              "fb_metrics": fb.metrics, "optim_metrics": opt.metrics, "seconds": time.time() - t0}
print(json.dumps(out)[:3000], flush=True)
# Fixed batch: 4 train problems x 8 rollouts from this client's sampler.
t1 = time.time()
samp = tc.save_weights_and_get_sampling_client(name="rx-b40a-preflight")
probs = ge.load("gsm8k_train.jsonl")[:4]
groups = rl.rollouts(samp, probs, pool)
datums, flat, cons = rl.build_datums(samp, groups, pool)
out["fixed_batch_seconds"] = time.time() - t1; out["n_datums"] = len(datums)
# Persist the fixture
fixture = [{"tokens": s["prompt"] + s["tokens"], "prompt_len": len(s["prompt"]), "samp_lp": s["logprobs"], "adv": s["adv"]} for s in flat]
blob = json.dumps({"fixture": fixture, "datums": [{"model_input": dd.model_input.to_ints(), **{k: v.to_numpy().tolist() for k, v in dd.loss_fn_inputs.items()}} for dd in datums]}).encode()
open("preflight_fixture.json", "wb").write(blob); out["fixture_sha256"] = hashlib.sha256(blob).hexdigest()
# Stock built-in importance_sampling on 1-D targets
stock = []
for dd, s in zip(datums, flat):
    inp = dd.loss_fn_inputs; N = inp["mask"].to_numpy().shape[0]
    stock.append(types.Datum(model_input=dd.model_input, loss_fn_inputs={
        "target_tokens": types.TensorData.from_numpy(inp["target_tokens"].to_numpy()[:, 0].astype(np.int64)),
        "logprobs": types.TensorData.from_numpy(inp["samp_lp"].to_numpy()),
        "advantages": types.TensorData.from_numpy(inp["adv"].to_numpy() * inp["mask"].to_numpy())}))
fb_stock = tc.forward_backward(stock, "importance_sampling").result()
stock_lp = [o["logprobs"].to_numpy() for o in fb_stock.loss_fn_outputs]
store0, store1 = [], []
fb0 = tc.forward_backward_custom(datums, rl.make_loss(0, store0)).result()
fb1 = tc.forward_backward_custom(datums, rl.make_loss(1, store1)).result()
out["stock_metrics"] = fb_stock.metrics; out["custom_c0_metrics"] = fb0.metrics; out["custom_c1_metrics"] = fb1.metrics
torch.save({"stock_lp": stock_lp, "lp_c0": store0, "lp_c1": store1}, "preflight_logprobs.pt")
# Baseline parity at the logprob interface: stock importance_sampling (tinker docs: loss = -sum(exp(lp - samp_lp) * adv))
# vs custom CENTER=0 through forward_backward_custom, using each path's own server logprobs.
import sys; sys.path.insert(0, "/home/ubuntu/repos/research-exchange")
from research_exchange.correctness import parity_errors
ref_l, act_l, ref_g, act_g, lp_ref, lp_act = 0.0, 0.0, [], [], [], []
for dd, slp, lp0 in zip(datums, stock_lp, store0):
    inp = {k: v.to_torch() for k, v in dd.loss_fn_inputs.items() if k not in ("target_tokens", "weights")}
    a = inp["adv"] * inp["mask"]
    x = torch.tensor(slp, dtype=torch.float32, requires_grad=True)
    l = -(torch.exp(x - inp["samp_lp"]) * a).sum(); l.backward(); ref_l += l.item(); ref_g.append(x.grad.clone())
    y = lp0.clone().requires_grad_(True)
    l2, _ = rl.make_loss(0)([dd], [y]); l2.backward(); act_l += l2.item(); act_g.append(y.grad[:, 0].clone())
    assert float(y.grad[:, 1:].abs().max()) == 0.0
    lp_ref.append(torch.tensor(slp)[inp["mask"] > 0]); lp_act.append(lp0[:, 0][inp["mask"] > 0])
out["correctness_server"] = {
    "stock_loss_sum_server_metric": fb_stock.metrics.get("loss:sum"), "custom_c0_loss_local": act_l, "stock_loss_local": ref_l,
    "baseline_loss": parity_errors([ref_l], [act_l]),
    "baseline_gradient": parity_errors(torch.cat(ref_g).tolist(), torch.cat(act_g).tolist()),
    "baseline_logprobs": parity_errors(torch.cat(lp_ref).tolist(), torch.cat(lp_act).tolist()),
    "n_tokens": int(sum(float(d.loss_fn_inputs["mask"].to_numpy().sum()) for d in datums))}
json.dump(out, open("preflight_result.json", "w"), indent=1, default=str)
print(json.dumps(out, default=str)[:4000], flush=True)