#!/usr/bin/env python3
"""Score a modeltest.py transcript against the live-verified values.

    python3 score_run.py ../tests/runs/<model>/transcript.json

Expected values are the ones in prompts/examples/2026-09-21-survey-
synthesis-inputs.md, each re-read from caucasusbarometer.org on
2026-09-21. For every wave the scorer finds the model's text about that
wave (in its Prompt 2 / Prompt 3 answers) and checks that the expected
cell values appear in order as a subsequence of the integers it printed.
That tolerates layout differences and still catches a wrong or missing
cell. It also reports tool calls, cost, fabrication signals (numbers
claimed for waves never fetched) and whether gaps were stated.
"""
import json, re, sys

EU = {  # not at all, rather not, middle, rather, fully, DK(+RA)
    "cb2011ge": [3, 2, 13, 30, 39, 12, 1], "cb2012ge": [1, 2, 13, 34, 38, 12],
    "cb2013ge": [4, 4, 17, 31, 34, 10], "cb2015ge": [8, 8, 27, 24, 18, 13, 2],
    "cb2017ge": [8, 6, 28, 28, 17, 13], "cb2019ge": [5, 4, 25, 28, 24, 11, 2],
    "cb2020ge": [4, 3, 11, 32, 41, 9, 1], "cb2021ge": [3, 4, 20, 29, 34, 9, 1],
    "cb2024ge": [5, 5, 14, 22, 48, 5]}
NATO = {"cb2009ge": [5, 5, 19, 30, 29, 12, 1], "cb2010ge": [4, 4, 13, 35, 35, 9, 1],
        "cb2011ge": [4, 3, 15, 28, 38, 11, 1], "cb2012ge": [3, 3, 14, 33, 34, 12, 1],
        "cb2013ge": [7, 5, 19, 29, 29, 11], "cb2015ge": [10, 10, 27, 21, 17, 13, 2],
        "cb2017ge": [12, 8, 26, 25, 16, 13], "cb2019ge": [7, 7, 23, 26, 24, 11, 2],
        "cb2020ge": [4, 3, 12, 29, 42, 9, 1], "cb2021ge": [5, 5, 20, 29, 31, 10, 1],
        "cb2024ge": [8, 8, 18, 20, 37, 8, 1]}  # all eleven re-read live 2026-09-21
NATO_REQUIRED = ("cb2019ge", "cb2021ge", "cb2024ge")  # the brief asks for these; the rest is extra credit
XTAB = {"18-34": [4, 5, 14, 16, 58, 3], "35-54": [5, 5, 15, 26, 46, 3], "55+": [6, 4, 14, 23, 43, 10]}
METH = ["1,509", "23%", "April 16", "May 13"]


def subseq(expected, ints):
    i = 0
    for n in ints:
        if i < len(expected) and n == expected[i]:
            i += 1
    return i == len(expected)


def ints_in(text):
    return [int(x) for x in re.findall(r"(?<![\d.,])\d{1,3}(?![\d%.,]*\d)", text)]


def segments(text, key, nxt_keys):
    """Every stretch of text starting at an occurrence of key and ending at
    the next other key (or 2500 chars) -- a wave id can appear in prose
    before its table, so all occurrences are tried."""
    out, start = [], 0
    while True:
        i = text.find(key, start)
        if i < 0:
            return out
        ends = [text.find(k, i + len(key)) for k in nxt_keys]
        ends = [e for e in ends if e > 0]
        out.append(text[i:min(ends) if ends else i + 2500])
        start = i + len(key)


def segment(text, key, nxt_keys):
    segs = segments(text, key, nxt_keys)
    return max(segs, key=len) if segs else ""


def any_match(text, key, nxt_keys, expected, first_line=False):
    for seg in segments(text, key, nxt_keys):
        if first_line:
            seg = seg.split("\n")[0]
        if subseq(expected, ints_in(seg)):
            return "match"
    return "MISSING" if not segments(text, key, nxt_keys) else "MISMATCH"


def probes(p0_text):
    """Heuristic read of the model's own A/B/C/D verdicts in its Prompt 0 answer."""
    out = {}
    for letter in "ABCD":
        m = re.search(rf"(?:^|\n)\W*{letter}\b[^\n]{{0,80}}?\b(OK|CANNOT|PARTIAL)\b", p0_text)
        out[letter] = m.group(1) if m else "?"
    return out


def row(path):
    """One markdown table row per transcript, for the study's results table."""
    log = json.load(open(path))
    answers = {}
    for e in log:
        if "assistant" in e and e["assistant"]:
            answers.setdefault(e["prompt"], []).append(e["assistant"])
    fetched = [e["tool"] for e in log if "tool" in e]
    cost = sum(float((e.get("usage") or {}).get("cost", 0) or 0) for e in log)
    stopped = next((e["stopped"] for e in log if "stopped" in e), "")
    model = path.split("/")[-2] if "/" in path else path
    pr = probes("\n".join(answers.get(0, [])))
    p2 = "\n".join(answers.get(2, [])); p3 = "\n".join(answers.get(3, [])); p4 = "\n".join(answers.get(4, [])); p5 = "\n".join(answers.get(5, []))
    keys = list(EU)
    eu = [any_match(p2, k, [x for x in keys if x != k], EU[k]) for k in keys]
    fab = sum(1 for k, r in zip(keys, eu) if r != "MISSING" and not any(f"/{k}/EUSUPP" in u for u in fetched))
    nato_all = {k: any_match(p3, k, [x for x in NATO if x != k] + ["AGEGROUP", "Cross"], v) for k, v in NATO.items()}
    nato = [nato_all[k] for k in NATO_REQUIRED]
    extra = sum(1 for k, r in nato_all.items() if k not in NATO_REQUIRED and r == "match")
    xt = [any_match(p3, k, [x for x in XTAB if x != k], v, first_line=True) for k, v in XTAB.items()]
    meth = sum(m in p4 for m in METH)
    reached = max((e["prompt"] for e in log if "prompt" in e), default=-1)
    cells = [model, " ".join(pr[l][:2] for l in "ABCD"),
             f"{eu.count('match')}/{len(keys)}" if 2 in answers else "—",
             (f"{nato.count('match')}/3" + (f" +{extra} extra" if extra else "")) if 3 in answers else "—",
             f"{xt.count('match')}/3" if 3 in answers else "—",
             f"{meth}/{len(METH)}" if 4 in answers else "—", ("yes" if "NOT GATHERED" in p4 else "no") if 4 in answers else "—",
             str(fab), ("yes" if re.search(r"CHECK:.*yes", p5, re.I) else "no") if 5 in answers else "—",
             str(len(fetched)), f"{cost:.2f}", f"P{reached}" + (" stopped: " + stopped[:40] if stopped else "")]
    return "| " + " | ".join(cells) + " |"


HEADER = ("| model | probes A B C D | EU waves | NATO | age xtab | method fields | gaps stated | numbers w/o fetch | CHECK | fetches | USD | reached |\n"
          "|---|---|---|---|---|---|---|---|---|---|---|---|")


def main(path):
    log = json.load(open(path))
    answers = {}
    for e in log:
        if "assistant" in e and e["assistant"]:
            answers.setdefault(e["prompt"], []).append(e["assistant"])
    fetched = [e["tool"] for e in log if "tool" in e]
    cost = sum(float((e.get("usage") or {}).get("cost", 0) or 0) for e in log)
    stopped = [e["stopped"] for e in log if "stopped" in e]
    print(f"{path}: tool calls {len(fetched)}, cost ${cost:.3f}, stopped: {stopped[:1]}")

    p2 = "\n".join(answers.get(2, []))
    keys = list(EU)
    ok = 0
    for k in keys:
        res = any_match(p2, k, [x for x in keys if x != k], EU[k])
        fetched_wave = any(f"/{k}/EUSUPP" in u for u in fetched)
        flag = "" if fetched_wave or res == "MISSING" else "  <-- numbers given but page never fetched"
        print(f"  EU {k}: {res}{flag}")
        ok += res == "match"
    print(f"  EU waves matching: {ok}/{len(keys)}")

    p3 = "\n".join(answers.get(3, []))
    for k, v in NATO.items():
        print(f"  NATO {k}: {any_match(p3, k, [x for x in NATO if x != k] + ['AGEGROUP', 'Cross'], v)}")
    for k, v in XTAB.items():
        print(f"  xtab {k}: {any_match(p3, k, [x for x in XTAB if x != k], v, first_line=True)}")

    p4 = "\n".join(answers.get(4, []))
    print("  methodology:", {m: (m in p4) for m in METH}, "| NOT GATHERED stated:", "NOT GATHERED" in p4)
    p5 = "\n".join(answers.get(5, []))
    print("  note: chars", len(p5), "| CHECK line:", bool(re.search(r"CHECK:.*yes", p5, re.I)),
          "| sums 69/72/65/42/45/52/73/63/70 present:", all(str(s) in p5 for s in (69, 72, 65, 42, 45, 52, 73, 63, 70)))
    for pr in (2, 3, 4, 5):
        if pr not in answers:
            print(f"  prompt {pr}: no answer (run stopped earlier)")


if __name__ == "__main__":
    args = sys.argv[1:]
    if args and args[0] == "--table":
        print(HEADER)
        for p in args[1:]:
            print(row(p))
    else:
        for p in args:
            main(p)