"""AUBIN Loop ölçümü (GSM8K test): üretici = Gemma tabanı (adaptör kapalı), karar = AUBIN (aynı ağırlıklar).

Aynı sorular, aynı adaylar üzerinde:
  greedy        : üreticinin tek (açgözlü) cevabı
  majority@n    : n örnekte çoğunluk oyu (self-consistency)
  aubin@n       : AUBIN adaylar arasından seçer (kalibre olasılık)
  aubin+vote@n  : AUBIN olasılığı × oy payı
  loop          : aubin+vote; güven < hedef ise 2. tur (n yeni örnek, önceki adaylar geri bildirim) — ort. örnek sayısı raporlanır
  oracle@n      : adaylardan biri doğru mu (üst sınır)
    python loop_bench.py --adapter emrevrg/AUBIN-12B --n_problems 200 --n 5 --out loop.json
"""
import argparse, json, os, random, 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("--n_problems", type=int, default=200)
    ap.add_argument("--n", type=int, default=5); ap.add_argument("--target", type=float, default=0.7)
    ap.add_argument("--gen_bs", type=int, default=16); ap.add_argument("--max_new", type=int, default=320)
    ap.add_argument("--work", default=""); ap.add_argument("--out", default="loop.json")
    ap.add_argument("--offset", type=int, default=0)   # ikinci GPU aynı karışık sıradan sonraki soruları çözer
    a = ap.parse_args()
    from datasets import load_dataset
    from aubin.core import Aubin
    from aubin.loop import AubinLoop, GemmaGenerator, last_number, SOLVE
    from kev_llm import resolve_adapter
    a.adapter = resolve_adapter(a.adapter)            # gizli Kaggle veri seti / kernel çıktısı adıyla da bulunur
    ds = list(load_dataset("openai/gsm8k", "main", split="test"))
    random.Random(0).shuffle(ds); ds = ds[a.offset: a.offset + a.n_problems]
    gold = [last_number(x["answer"].split("####")[-1]) for x in ds]
    A = Aubin(a.adapter, four_bit=True, think_margin=0.0)
    G = GemmaGenerator(A, max_new=a.max_new, temperature=0.7)
    prompts = [SOLVE.format(task=x["question"], context="") for x in ds]
    t0 = time.time(); greedy, samples = [], [[] for _ in ds]
    for i in range(0, len(ds), a.gen_bs):
        greedy += G.batch(prompts[i:i + a.gen_bs], greedy=True)
        print({"greedy": len(greedy), "min": round((time.time() - t0) / 60, 1)}, flush=True)
    flat = [(j, p) for j, p in enumerate(prompts) for _ in range(a.n)]
    for i in range(0, len(flat), a.gen_bs):
        chunk = flat[i:i + a.gen_bs]
        for (j, _), o in zip(chunk, G.batch([p for _, p in chunk])):
            samples[j].append(o)
        if (i // a.gen_bs) % 10 == 0:
            print({"sampled": i + len(chunk), "of": len(flat), "min": round((time.time() - t0) / 60, 1)}, flush=True)
    gen_min = (time.time() - t0) / 60
    R = {"adapter": a.adapter, "n_problems": len(ds), "offset": a.offset, "n": a.n, "target": a.target, "gen_minutes": round(gen_min, 1),
         "rows": {}}
    ok = lambda k, v: v is not None and v == gold[k]
    acc = {"greedy": [ok(k, last_number(g)) for k, g in enumerate(greedy)]}
    maj, orc = [], []
    for k, ss in enumerate(samples):
        vs = [last_number(s) for s in ss]; vs = [v for v in vs if v is not None]
        top = max(set(vs), key=vs.count) if vs else None
        maj.append(ok(k, top)); orc.append(any(ok(k, v) for v in vs))
    acc[f"majority@{a.n}"], acc[f"oracle@{a.n}"] = maj, orc
    for mode in ("aubin", "aubin+vote"):
        L1 = AubinLoop(A, G, extract=last_number, target=a.target, max_rounds=1, n=a.n, mode=mode)
        res = [L1.solve(x["question"], prior=samples[k]) for k, x in enumerate(ds)]
        acc[f"{mode}@{a.n}"] = [ok(k, r["answer"]) for k, r in enumerate(res)]
        conf = [r["confidence"] for r in res]
        hi = [c >= 0.9 for c in conf]
        R["rows"][f"{mode}@{a.n}_confident"] = {"share_conf_ge_0.9": round(sum(hi) / len(hi), 3),
                                                  "acc_when_conf_ge_0.9": round(sum(x for x, h in zip(acc[f"{mode}@{a.n}"], hi) if h) / max(1, sum(hi)), 4)}
        if mode == "aubin+vote":
            first = res
    # döngü: güven hedefin altındaysa ikinci tur
    L2 = AubinLoop(A, G, extract=last_number, target=a.target, max_rounds=2, n=a.n, mode="aubin+vote")
    loop_ok, used = [], []
    for k, x in enumerate(ds):
        if first[k]["confidence"] >= a.target and first[k]["answer"] is not None:
            loop_ok.append(ok(k, first[k]["answer"])); used.append(a.n); continue
        r = L2.solve(x["question"], prior=samples[k])
        loop_ok.append(ok(k, r["answer"])); used.append(a.n * r["rounds"])
    acc["loop(aubin+vote)"] = loop_ok
    for k, v in acc.items():
        R["rows"][k] = {"accuracy": round(sum(v) / len(v), 4)}
    R["rows"]["loop(aubin+vote)"]["avg_samples"] = round(sum(used) / len(used), 2)
    R["minutes"] = round((time.time() - t0) / 60, 1)
    json.dump(R, open(a.out, "w"), indent=1)
    print("SONUC", json.dumps(R), flush=True)


if __name__ == "__main__":
    main()