"""AUBIN vs Jev — Jev'in KENDİ raporladığı küme üzerinde (Kev development, variant == "clean").

Jev'in yayımlanan skorları (Kev repo runs/jev-*-v4/report.json → "clean") development bölümünün varyantsız
sorularıdır (transfer 656 / decision 1264). Varyantlar (none_absent, none_present, permuted) ayrı raporlanır.
Bu betik aynı soruları _meta.variant ile seçer ve ürünün (aubin.core) yaptığı hesabın aynısını uygular:

    üye log-olas. = log_softmax(lp / T_üye); eminsiz sorularda (hızlı yolda ilk iki farkı < marj) düşünen yol
    think_mix oranında karıştırılır; üyeler ağırlıkla toplanır; birleşim / T_son → softmax.

    python jev_compare.py out.json --spec spec.json
spec: {"rows": {ad: {"members": [{"run": think-run.json | {bölüm: dosya}, "T": 3.45, "w": 0.6, "think_mix": 0.5}], "T": 0.8}},
       "jev": {"kev_transfer_dev": jev-transfer-v4.json, "kev_dev": jev-decision-v4.json}}
"""
import argparse, json, math, os
import numpy as np

SUITE_FILE = {"kev_transfer_dev": "v4/transfer-v4/development.jsonl", "kev_dev": "v7/decision-v7/development.jsonl",
              "kev_transfer_test": "v4/transfer-v4/test.jsonl", "kev_test": "v7/decision-v7/test.jsonl"}


def variants(tag, work):
    from huggingface_hub import hf_hub_download
    p = hf_hub_download("jaredpalmer/kev-suites", SUITE_FILE[tag], repo_type="dataset", local_dir=work)
    return [r["_meta"].get("variant", "clean") for l in open(p, encoding="utf-8") for r in [json.loads(l)] for _ in r["questions"]]


def load_run(path):
    """Koşu JSON'u; .gz ise sıkıştırılmış okunur (disk dar olduğunda sonuçlar gzip'li saklanır)."""
    import gzip
    op = gzip.open if str(path).endswith(".gz") else open
    with op(path, "rt", encoding="utf-8") as f:
        return json.load(f)


def lsm(x):
    x = np.asarray(x, dtype=np.float64)
    return x - np.logaddexp.reduce(x)


def metrics(P, Y):
    acc = br = nll = 0.0; conf = []
    for p, y in zip(P, Y):
        oh = np.zeros(len(p)); oh[y] = 1
        acc += int(p.argmax()) == y; br += float(((p - oh) ** 2).sum()); nll -= math.log(max(p[y], 1e-12))
        conf.append((float(p.max()), int(p.argmax()) == y))
    n = len(Y); ece = 0.0
    for j in range(15):
        b = [c for c in conf if j / 15 < c[0] <= (j + 1) / 15]
        if b:
            ece += len(b) / n * abs(np.mean([c[0] for c in b]) - np.mean([c[1] for c in b]))
    return {"n": n, "accuracy": round(acc / n, 4), "brier": round(br / n, 4), "ece15": round(ece, 4), "nll": round(nll / n, 4)}


def sources(tag, work):
    from huggingface_hub import hf_hub_download
    p = hf_hub_download("jaredpalmer/kev-suites", SUITE_FILE[tag], repo_type="dataset", local_dir=work)
    return [q.get("src") or r["_meta"]["source"] for l in open(p, encoding="utf-8") for r in [json.loads(l)] for q in r["questions"].values()]


def combine(spec, tag, V, margin=4.0, return_probs=False):
    idx = [i for i, v in enumerate(V) if v == "clean"]
    acc, Y = None, None
    for m in spec["members"]:
        d = load_run(m["run"]).get("runs", {}).get(tag)
        if not d or "items" not in d:
            return None
        items, fast = d["items"], d.get("fast_items")
        if Y is None:
            Y = [items[i]["y"] for i in idx]
        out = []
        for i in idx:
            f = lsm(np.array(fast[i] if fast else items[i]["lp"]))
            z = lsm(f / m["T"])
            srt = np.sort(f)[::-1]
            if fast and m.get("think_mix", 0) > 0 and 1 < len(srt) <= 26 and srt[0] - srt[1] < margin:
                z = (1 - m["think_mix"]) * z + m["think_mix"] * lsm(lsm(np.array(items[i]["lp"])) / m["T"])
            out.append(m["w"] * z)
        acc = out if acc is None else [a + b for a, b in zip(acc, out)]
    P = [np.exp(lsm(a / spec["T"])) for a in acc]
    if return_probs:
        return P, Y, idx
    return metrics(P, Y)


def extra(P, Y, src, jev_report):
    """Kural/politika (composition_*, contrastive_*) hata sayısı ve 'emin ama yanlış' (güven >= 0.9) oranı, Jev ile yan yana."""
    rule = lambda s: s.startswith(("composition", "contrastive"))
    err = [int(int(p.argmax()) != y) for p, y in zip(P, Y)]
    conf = [float(p.max()) for p in P]
    jt, jc = jev_report["tasks"], jev_report["clean"]
    return {"rule_policy_n": sum(1 for s in src if rule(s)),
            "rule_policy_errors": sum(e for e, s in zip(err, src) if rule(s)),
            "jev_rule_policy_errors": round(sum(v["n"] * (1 - v["acc"]) for k, v in jt.items() if rule(k))),
            "confident_error_rate": round(sum(e for e, c in zip(err, conf) if c >= 0.9) / len(err), 4),
            "jev_confident_error_rate": round(jc["confident_error_rate"], 4),
            "errors": sum(err), "jev_errors": round(jc["n"] * (1 - jc["acc"]))}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("out"); ap.add_argument("--spec", required=True)
    ap.add_argument("--work", default=os.path.join(os.environ.get("TEMP", "/tmp"), "kev_suites"))
    a = ap.parse_args()
    spec = json.load(open(a.spec, encoding="utf-8"))
    res = {"spec": spec, "subset": "Kev development, variant == clean (the set Jev's published scores use)", "rows": {}}
    for tag, jf in spec["jev"].items():
        V = variants(tag, a.work)
        j = json.load(open(jf, encoding="utf-8"))["clean"]
        res["rows"].setdefault("Jev (hosted, published report)", {})[tag] = {
            "n": j["n"], "accuracy": round(j["acc"], 4), "brier": round(j["brier"], 4), "ece15": round(j["ece"], 4), "nll": round(j["nll"], 4)}
        for name, row in spec["rows"].items():
            ms = [dict(m, run=m["run"][tag] if isinstance(m["run"], dict) else m["run"]) for m in row["members"]]
            r = combine({"members": ms, "T": row["T"]}, tag, V)
            if r:
                res["rows"].setdefault(name, {})[tag] = {k: (round(float(v), 4) if k != "n" else v) for k, v in r.items()}
                P, Y, idx = combine({"members": ms, "T": row["T"]}, tag, V, return_probs=True)
                src = sources(tag, a.work)
                res["rows"][name][tag]["vs_jev"] = extra(P, Y, [src[i] for i in idx], json.load(open(jf, encoding="utf-8")))
    json.dump(res, open(a.out, "w", encoding="utf-8"), indent=1)
    for name, v in res["rows"].items():
        print(name, {t: (x["accuracy"], x["brier"], x["ece15"], x["nll"], x.get("vs_jev")) for t, x in v.items()})


if __name__ == "__main__":
    main()