"""AUBIN hız: durum KV önbelleği (aynı duruma sorulan sorularda durum bir kez işlenir) — doğruluk eşliği + soru-başı gecikme.

Aynı model, aynı sorular (Kev dev durumları, çok sorulu), önbellek KAPALI ↔ AÇIK:
  - fark: log-olasılıklarda en büyük mutlak fark, argmax uyuşması
  - gecikme: soru başına ms (tek soru, grup boyutu 1 — Jev'in API gecikmesiyle kıyaslanabilir ölçü)
    python bench_cache.py --adapter emrevrg/AUBIN-12B --n_cases 60 --out cache.json
"""
import argparse, json, os, sys, time
import torch

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--adapter", default="emrevrg/AUBIN-12B"); ap.add_argument("--no4bit", action="store_true")
    ap.add_argument("--device_map", default=None); ap.add_argument("--n_cases", type=int, default=60)
    ap.add_argument("--suite", default="kev_dev"); ap.add_argument("--work", default="/tmp/kev"); ap.add_argument("--out", default="cache.json")
    a = ap.parse_args()
    import kevdata
    from aubin.core import Aubin
    from kev_llm import resolve_adapter
    a.adapter = resolve_adapter(a.adapter)
    K = kevdata.load(a.work)
    cases = [c for c in K[a.suite] if len(c["questions"]) >= 2][: a.n_cases]
    A = Aubin(a.adapter, four_bit=not a.no4bit, device_map=a.device_map, think_margin=0.0)
    R = {"adapter": a.adapter, "four_bit": not a.no4bit, "gpu": torch.cuda.get_device_name(0), "suite": a.suite,
         "cases": len(cases), "questions": sum(len(c["questions"]) for c in cases)}
    A.cache_state = False; A.logprobs(cases[0]["state"], dict(list(cases[0]["questions"].items())[:1]))   # ısınma
    res = {}
    for mode in (False, True):
        A.cache_state, A._pc = mode, None
        torch.cuda.synchronize(); t0 = time.time(); lps = []
        for c in cases:
            lp = A.logprobs(c["state"], c["questions"])
            lps.append({k: v[1] for k, v in lp.items()})
        torch.cuda.synchronize(); dt = time.time() - t0
        res[mode] = lps
        R["cache_on" if mode else "cache_off"] = {"ms_per_question": round(1000 * dt / R["questions"], 1), "seconds": round(dt, 1)}
        print(mode, R["cache_on" if mode else "cache_off"], flush=True)
    diffs, agree, n = [], 0, 0
    for x, y in zip(res[False], res[True]):
        for k in x:
            diffs.append(float((x[k] - y[k]).abs().max())); agree += int(x[k].argmax() == y[k].argmax()); n += 1
    R["max_abs_logprob_diff"] = round(max(diffs), 4); R["mean_abs_logprob_diff"] = round(sum(diffs) / len(diffs), 5)
    R["argmax_agreement"] = round(agree / n, 4)
    R["speedup"] = round(R["cache_off"]["ms_per_question"] / R["cache_on"]["ms_per_question"], 2)
    json.dump(R, open(a.out, "w"), indent=1)
    print("SONUC", json.dumps(R), flush=True)


if __name__ == "__main__":
    main()