"""Run THX-01 on the exact nibzard decision-model-benchmark suites (hash-verified), same state & criteria as Jev."""
import json, sys, time, collections
import numpy as np
sys.path.insert(0, ".")
import thx01
S = "data/nibzard_suites/"
INS = ("Choose the single best option for the decision state. You must commit to a single option even if the state does not "
       "determine the answer; in that case choose the least unreasonable option and report low confidence. Answer confidence "
       "separately. The confidence is your probability that your chosen option is correct.")


def ece(conf, cor, bins=10):
    conf, cor = np.asarray(conf, float), np.asarray(cor, float)
    idx = np.digitize(conf, np.linspace(0, 1, bins + 1)[1:-1], right=True)
    return float(sum((idx == b).sum() / len(conf) * abs(cor[idx == b].mean() - conf[idx == b].mean()) for b in range(bins) if (idx == b).any()))


MP = sys.argv[1] if len(sys.argv) > 1 else "."
ag = thx01.load(MP, device="cuda")
out = {}
for suite in ["s1_intent77", "s2_spam", "s3_cardinality", "s4_order", "s5_confidence"]:
    items = [json.loads(l) for l in open(S + suite + ".jsonl")]
    qs = [{"decision": {"type": "choice", "instructions": INS, "criteria": {o: o for o in it["options"]}}} for it in items]
    t0 = time.perf_counter()
    res = ag.predict_batch([it["state"].strip() for it in items], qs, batch_size=32)
    wall = time.perf_counter() - t0
    ch, cf = [], []
    for it, r in zip(items, res):
        a = r["answers"]["decision"]; ch.append(it["options"].index(a["choice"])); cf.append(float(a["probabilities"][a["choice"]]))
    gold = [int(it["gold_index"]) for it in items]
    g = [i for i, x in enumerate(gold) if x >= 0]
    cor = [ch[i] == gold[i] for i in g]
    m = {"n": len(items), "accuracy": float(np.mean(cor)), "ece": ece([cf[i] for i in g], cor),
         "brier": float(np.mean([(cf[i] - float(ch[i] == gold[i])) ** 2 for i in g])), "ms_per_item": wall / len(items) * 1000}
    if suite == "s4_order":
        by = collections.defaultdict(list)
        for it, c in zip(items, ch):
            by[it["base_item_id"]].append(it["options"][c])
        m["flip_rate"] = sum(len(set(v)) > 1 for v in by.values()) / len(by)
    if suite == "s5_confidence":
        ng = [cf[i] for i, x in enumerate(gold) if x < 0]
        m["admits_ignorance"] = float(np.mean([c <= 0.5 for c in ng])); m["mean_conf_no_good"] = float(np.mean(ng))
    if suite == "s3_cardinality":
        byn = collections.defaultdict(list)
        for it, c, gg in zip(items, ch, gold):
            byn[json.loads(it["meta"].replace("'", '"'))["N"] if isinstance(it["meta"], str) else it["meta"]["N"]].append(c == gg)
        m["by_N"] = {int(k): float(np.mean(v)) for k, v in sorted(byn.items())}
    out[suite] = m
    print(suite, {k: (round(v, 4) if isinstance(v, float) else v) for k, v in m.items()}, flush=True)
json.dump(out, open(sys.argv[2] if len(sys.argv) > 2 else "./results/jev_suites_thx01.json", "w"), indent=1)