"""Evaluate the INT4 package against the original FP8 model's outputs (ref.json) under a GPU memory cap.
usage: python evaluate.py <model_dir> <cap_gib> <out.json>"""
import json, math, sys, time
import torch

path, cap, out_path = sys.argv[1], float(sys.argv[2]), sys.argv[3]
sys.path.insert(0, path); sys.path.insert(0, __import__("os").path.dirname(__file__))
total = torch.cuda.get_device_properties(0).total_memory
torch.cuda.set_per_process_memory_fraction(cap * 2**30 / total)          # hard cap: allocations beyond it OOM
from prism_int4 import Prism
from suite import SUITE

t = time.time()
model = Prism.from_pretrained(path, gpu_memory_gb=cap)
load_s = time.time() - t
ref = json.load(open(__import__("os").path.join(__import__("os").path.dirname(__file__), "eval_reference_fp8.json")))

# warm-up (Triton compile), then timed suite, one request per call
model.decide(**dict(zip(("state", "questions"), SUITE[0])))
torch.cuda.synchronize(); torch.cuda.reset_peak_memory_stats()
times, res = [], []
for s, q in SUITE:
    b0 = model.streamer.bytes_streamed; t = time.time()
    res.append(model.decide(state=s, questions=q)); torch.cuda.synchronize()
    times.append((time.time() - t, model.streamer.bytes_streamed - b0))
peak = torch.cuda.max_memory_allocated() / 2**30

# batched: all requests in one call
t = time.time(); res_b = model.decide_batch(SUITE); torch.cuda.synchronize(); batch_s = time.time() - t


def dist(a):
    if a["type"] == "yes_no":
        return {"no": 1 - a["yes"], "yes": a["yes"]}
    return a["probabilities"]


rows = []
for r_ref, r_q in zip(ref, res):
    for name in r_ref:
        p, qd = dist(r_ref[name]), dist(r_q[name])
        keys = list(p)
        kl = sum(p[k] * math.log(max(p[k], 1e-12) / max(qd[k], 1e-12)) for k in keys)
        rows.append(dict(top_agree=max(keys, key=p.get) == max(keys, key=qd.get),
                         max_abs=max(abs(p[k] - qd[k]) for k in keys), kl=kl, type=r_ref[name]["type"],
                         ref_top=max(p.values())))
n = len(rows)
batch_max_diff = max(abs(dist(a[k])[o] - dist(b[k])[o]) for ra, rb in zip(res, res_b) for k in ra
                     for a, b in [(ra, rb)] for o in dist(ra[k]))

# exception recovery: fail mid-forward inside a streamed layer, then check the next answer is unchanged
streamed = model.streamer.order
recovery = None
if len(streamed) > 2:
    mod = model.backbone.layers[streamed[2]].mlp
    orig = mod._moe_dispatch
    def boom(*a, **k):
        raise RuntimeError("injected failure")
    mod._moe_dispatch = boom
    try:
        model.decide(state=SUITE[1][0], questions=SUITE[1][1])
    except RuntimeError:
        pass
    mod._moe_dispatch = orig
    again = model.decide(state=SUITE[0][0], questions=SUITE[0][1])
    recovery = max(abs(dist(again[k])[o] - dist(res[0][k])[o]) for k in again for o in dist(again[k]))

summary = dict(
    cap_gib=cap, load_seconds=round(load_s), peak_gpu_gib=round(peak, 2),
    resident_layers=len(model.streamer.gpu), streamed_layers=len(streamed),
    questions=n, top1_agreement=sum(r["top_agree"] for r in rows) / n,
    mean_max_abs_prob_diff=sum(r["max_abs"] for r in rows) / n, worst_max_abs_prob_diff=max(r["max_abs"] for r in rows),
    mean_kl=sum(r["kl"] for r in rows) / n, worst_kl=max(r["kl"] for r in rows),
    top1_disagreements=[dict(r) for r in rows if not r["top_agree"]],
    median_request_seconds=sorted(x[0] for x in times)[len(times) // 2],
    gb_streamed_per_forward=round(sorted(x[1] for x in times)[len(times) // 2] / 1e9, 1),
    batched_all_requests_seconds=round(batch_s, 2), batched_vs_single_max_diff=batch_max_diff,
    recovery_after_injected_failure_max_diff=recovery)
json.dump(dict(summary=summary, results=res), open(out_path, "w"), ensure_ascii=False, indent=1)
print(json.dumps(summary, indent=1))