#!/usr/bin/env python3
"""Qwen (or any Ollama API model) on the EXACT released transitive-binding probe (name suites), chat surface,
first-choice accuracy by depth 1-5 — directly comparable to ConeML's heldout_transitive_probe numbers."""
import json, re, random, argparse, urllib.request

NAMES = ["Isabel", "Jonas", "Keira", "Liam", "Maya", "Noah", "Olivia", "Priya", "Quinn", "Rosa", "Sofia", "Theo"]
SUITES = {
    "sft_template_heldout_names": dict(rel="is taller than", hi="tallest", lo="shortest",
        hq="{c} Of all of them, the tallest is who? Return only the name.",
        lq="{c} Of all of them, the shortest is who? Return only the name."),
    "unseen_query_heldout_names": dict(rel="is taller than", hi="tallest", lo="shortest",
        hq="Given these facts: {c} Which person is highest in the height order? Answer with only the name.",
        lq="Given these facts: {c} Which person is lowest in the height order? Answer with only the name."),
    "older_relation_heldout_names": dict(rel="is older than", hi="oldest", lo="youngest",
        hq="{c} Which person is oldest? Return only the name.",
        lq="{c} Which person is youngest? Return only the name."),
}


def chain(items, rel): return " ".join(f"{items[i]} {rel} {items[i+1]}." for i in range(len(items)-1))


def first_choice(g, pool):
    t = g.lower(); hits = []
    for c in sorted(pool, key=len, reverse=True):
        m = re.search(rf"(?<![a-z]){re.escape(c.lower())}(?![a-z])", t)
        if m: hits.append((m.start(), c))
    return sorted(hits)[0][1] if hits else ""


def ask(host, model, prompt):
    body = json.dumps({"model": model, "messages": [{"role": "user", "content": prompt}], "stream": False,
                       "options": {"temperature": 0, "num_predict": 16}}).encode()
    req = urllib.request.Request(host.rstrip("/")+"/api/chat", data=body, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=60) as r:
        return json.loads(r.read())["message"]["content"]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True); ap.add_argument("--host", default="http://127.0.0.1:11434")
    ap.add_argument("--out", required=True); ap.add_argument("--n", type=int, default=40)
    a = ap.parse_args()
    res = {}
    for sname, s in SUITES.items():
        res[sname] = {}
        for depth in (1, 2, 3, 4, 5):
            rng = random.Random(1009 + depth * 7); ok = n = 0
            for _ in range(a.n // 2):
                items = rng.sample(NAMES, depth + 1); c = chain(items, s["rel"])
                for q, gold in [(s["hq"].format(c=c), items[0]), (s["lq"].format(c=c), items[-1])]:
                    try: g = ask(a.host, a.model, q)
                    except Exception: g = ""
                    n += 1; ok += int(first_choice(g, NAMES) == gold)
            res[sname][f"d{depth}"] = round(100 * ok / max(n, 1))
            print(f"  {sname} d{depth}: {res[sname][f'd{depth}']}% (n={n})", flush=True)
    json.dump({"model": a.model, "suites": res}, open(a.out, "w"), indent=1)
    print("\n=== %s transitive first-choice by depth ===" % a.model)
    for sn, dd in res.items(): print(f"{sn:32s} " + " ".join(f"{dd[f'd{d}']:>4}" for d in range(1, 6)))
    print("->", a.out)


if __name__ == "__main__":
    main()