# E18: reversible coordinate preconditioning on MATH-500 (Sol round-2 idea; prior-art verdict CLOSE; NO novelty claim).
# Three one-solution arms at matched compute (one greedy 14B solve, 2048 cap, same server as E17):
#   I = identity (original statement) -> reused from E17 greedy results
#   N = fixed conventional normalization (control);  S = min-J selector.
# Items where an arm's prompt equals the original reuse E17; identical N/S prompts are generated once.
# Also reruns the original prompt on every eligible item to measure run-to-run nondeterminism (I_rerun).
import json, time, requests, sys
URL = "http://127.0.0.1:8094"; OUT = "results/quality/e18_precond.jsonl"
plan = json.load(open("bench/precond/plan_math500.json", encoding="utf-8"))
ds = [json.loads(l) for l in open("bench/precond/math500.jsonl", encoding="utf-8")]
SYS = "<|im_start|>system\nPlease reason step by step, and put your final answer within \\boxed{}.<|im_end|>\n<|im_start|>user\n"
def gen(user):
    p = SYS + user + "<|im_end|>\n<|im_start|>assistant\n"; t = time.time()
    j = requests.post(URL + "/completion", json={"prompt": p, "n_predict": 2048, "cache_prompt": False, "temperature": 0, "top_k": 1, "seed": 0}, timeout=3600).json()
    return {"text": j.get("content", ""), "timings": j.get("timings", {}), "wall": round(time.time() - t, 2)}
with open(OUT, "w", encoding="utf-8") as f:
    for r in plan:
        k = r["item"]; jobs = {}
        if r["prompt_N"] or r["prompt_S"]: jobs["I_rerun"] = ds[k]["problem"]
        if r["prompt_N"]: jobs["N"] = r["prompt_N"]
        if r["prompt_S"] and r["prompt_S"] != r["prompt_N"]: jobs["S"] = r["prompt_S"]
        for arm, pr in jobs.items():
            o = gen(pr); o.update({"item": k, "arm": arm, "prompt_user": pr})
            f.write(json.dumps(o) + "\n"); f.flush()
            print(k, arm, o["timings"].get("predicted_n"), o["wall"], flush=True)
print("E18PYDONE", flush=True)