#!/usr/bin/env python3
"""parse_v2.py -- the cost-of-intelligence instrument's strict parser (parser v2), frozen by PREREG-S0.md.

    python3 parse_v2.py --selftest            # every registered case; exit 1 on any mismatch
    python3 parse_v2.py --selftest --table    # the same cases as the markdown table PREREG-S0.md carries
    python3 parse_v2.py --census ROWS.jsonl --set c [--stop-field done_reason]
                                              # re-classify one row file: class counts, strict / lenient / v1 correct

WHAT IT DECIDES. One reply's visible text (never a thinking channel) -> a format class, a STRICT choice
(the only one that scores), a LENIENT choice (printed beside, never in a ratio) and run.py's v1 choice
(kept for continuity). The rule is PREREG-S0.md section 3; this file is its reference implementation
and PREREG-S0.md records this file's sha256. A change here is a new parser version and an amendment.

THE RULE, in order (set_id names the option letters and the names a reply may use for each option):
  0. an EMPTY reply is `budget_exhausted` when the server's stop reason is "length" (the budget ran out
     before any visible answer), else `unparsed`; both score wrong, neither is ever turned into a choice.
  1. the OPENING letter: after leading whitespace and markdown decoration (* _ ` # >), an UPPERCASE option
     letter, optionally in [ ] or ( ), followed by end-of-text or a separator (whitespace ] ) . : , ; - * _ `).
     A lowercase letter is never a choice.
       1a. the English ARTICLE: an unbracketed "A" (no bracket on EITHER side: "A)" and "A]" are list
           markers, not the article) followed by exactly one space and a lowercase word on the same line,
           where that word neither starts an option's name nor is a choice-marker (is, was, would, seems,
           because, since, as, being, sounds, fits, matches, appears, then, so) -> `article_A`
           (no strict choice; lenient = the one option named anywhere, if exactly one).
       1b. the LABEL after the letter: the rest of the SAME LINE (up to 64 characters). If it names ANY
           option other than the letter's (with or without the letter's own), or states a different option
           letter ([X], (X), **X**, "answer: X") -> `contradictory` (no choice). If it names only the
           letter's own option -> `letter+name`. Otherwise -> `clean`. Text on later lines is explanation,
           not a label.
  2. no opening letter, doorman letter recast only: a JSON `"pass": true|false` -> `json_verdict`
     (lenient admit / refuse; no strict choice).
  3. STATED letters later in the reply: [X], (X), **X**, "answer/option (is)[:] X", or a final line that is
     exactly one letter (brackets, bold, a period allowed). Exactly one distinct option letter stated ->
     `late_letter` (lenient = it); two or more -> `contradictory`.
  4. NAMES: option names anywhere (word-bounded; longest alias first, each match consumed). Exactly one
     option named -> `name_only` (lenient = it); two or more -> `contradictory`.
  5. an option letter appears only inside words (v1 would have taken it) -> `word_letter`.
  6. otherwise `unparsed`.
STRICT correct  = class in {clean, letter+name} and the strict choice == label.
LENIENT correct = the lenient choice == label (never a headline, never in a ratio).
v1              = run.py's rule: the first option-letter CHARACTER anywhere; an unparsed reply becomes
                  option 0 (that is how run.py scored it; `correct_v1` reproduces run.py's count).
`template_suspect` flags a reply carrying a chat-template control string (</s>, <|im_start|>, ### Human, USER:).

The native doorman contract (set id doorman36-json) is scored by `json_verdict_native()`: the first
top-level JSON object with a boolean `pass` (true -> admit, false -> refuse); anything else is unparsed.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import sys
from dataclasses import dataclass

PARSER_VERSION = "v2"
LETTERS = [chr(65 + i) for i in range(26)]

#: The option letters and the names a reply may use for each option, per set id (PREREG-S0.md section 3).
#: Keys are the label strings the kits carry; aliases are matched word-bounded, case-insensitively.
SETS: dict[str, dict] = {
    "c": {"options": ["darwin", "einstein", "hypatia", "ibn_sina", "sagan", "socrates"],
          "names": {"darwin": ["charles darwin", "darwin", "charles"],
                    "einstein": ["albert einstein", "einstein", "albert"],
                    "hypatia": ["hypatia"],
                    "ibn_sina": ["ibn sina", "ibn_sina", "avicenna"],
                    "sagan": ["carl sagan", "sagan", "carl"],
                    "socrates": ["socrates"]}},
    "d48": {"options": ["admit", "refuse"],
            "names": {"admit": ["admit", "admitted", "hears the line", "pass"],
                      "refuse": ["refuse", "refused", "does not hear", "reject"]},
            "json_pass": True},
    "a": {"options": ["yes", "no"], "names": {"yes": ["yes"], "no": ["no"]}},
    "f40": {"options": ["yes", "no"], "names": {"yes": ["yes"], "no": ["no"]}},
    "j": {"options": ["grounded", "not_grounded", "off_page"],
          "names": {"grounded": ["grounded"],
                    "not_grounded": ["not grounded", "not_grounded"],
                    "off_page": ["off page", "off_page", "off-page"]}},
}
SET_ALIASES = {"d": "d48", "f": "f40"}   # the lanes' one-letter task names

ARTICLE_EXCEPTIONS = {"is", "was", "would", "seems", "because", "since", "as", "being", "sounds",
                      "fits", "matches", "appears", "then", "so"}
TEMPLATE_STRINGS = ("</s>", "<|im_start|>", "### Human", "USER:")
LABEL_CHARS = 64

_OPEN = re.compile(r"^([\[\(])?([A-Za-z])([\]\)])?(?=$|[\s\]\).:,;\-*_`])")
_STATED = re.compile(
    r"\[([A-Z])\]|\(([A-Z])\)|\*\*\s*\[?([A-Z])\]?\s*\*\*"
    r"|(?i:answer|option)\s*(?:is)?\s*[:\-]?\s*\**\s*\[?([A-Z])\]?(?![A-Za-z])")
_FINAL_LINE = re.compile(r"^[\[\(*_`\s]*([A-Z])[\]\)*_`\s]*\.?\s*$")
_JSON_PASS = re.compile(r'"pass"\s*:\s*(true|false)')


@dataclass(frozen=True)
class Parse:
    format_class: str
    choice_strict: str | None
    choice_lenient: str | None
    choice_v1: str | None      # run.py's choice: an unparsed reply becomes option 0, as run.py scored it
    parsed_v1: bool
    template_suspect: bool
    parser_version: str = PARSER_VERSION

    @property
    def strict_scorable(self) -> bool:
        return self.format_class in ("clean", "letter+name")


def set_spec(set_id: str) -> dict:
    return SETS[SET_ALIASES.get(set_id, set_id)]


def _named_in(text: str, spec: dict) -> list[str]:
    """Options named in `text`, in order of first appearance; longest alias first, each match consumed."""
    low = text.lower()
    pairs = sorted(((a, k) for k, al in spec["names"].items() for a in al), key=lambda p: -len(p[0]))
    hits: dict[str, int] = {}
    for alias, key in pairs:
        for m in re.finditer(r"(?<![a-z0-9_])" + re.escape(alias) + r"(?![a-z0-9_])", low):
            if key not in hits or m.start() < hits[key]:
                hits[key] = m.start()
            low = low[:m.start()] + " " * (m.end() - m.start()) + low[m.end():]
    return [k for k, _ in sorted(hits.items(), key=lambda kv: kv[1])]


def choice_v1(text: str, n: int) -> tuple[str | None, bool]:
    """run.py's rule (jev-2026-09-21/run.py, `decide()`): the first option-letter character anywhere in the
    reply; no letter -> uniform probabilities -> `max()` returns index 0, so option 0 is the recorded choice."""
    for ch in text or "":
        if ch in LETTERS[:n]:
            return ch, True
    return LETTERS[0], False


def classify(text: str, set_id: str, stop_reason: str | None = None) -> Parse:
    spec = set_spec(set_id)
    opts = spec["options"]
    n = len(opts)
    letters = LETTERS[:n]
    raw = text or ""
    v1_letter, v1_parsed = choice_v1(raw, n)
    v1 = opts[letters.index(v1_letter)]
    suspect = any(t in raw for t in TEMPLATE_STRINGS)

    def mk(cls, strict, lenient):
        return Parse(cls, strict, lenient, v1, v1_parsed, suspect)

    if not raw.strip():
        return mk("budget_exhausted" if stop_reason == "length" else "unparsed", None, None)

    s = raw.lstrip()
    s2 = s.lstrip("*_`#> \n\t")
    named_all = _named_in(s, spec)
    lone = named_all[0] if len(named_all) == 1 else None

    m = _OPEN.match(s2)
    if m and m.group(2) in letters:
        L = m.group(2)
        pick = opts[letters.index(L)]
        bracketed = bool(m.group(1) or m.group(3))   # a bracket on either side: "[A]", "(A)", "A)", "A]"
        after = s2[m.end():]
        if L == "A" and not bracketed:
            am = re.match(r"^ ([a-z][a-z'\-]*)", after)
            if am:
                word = am.group(1)
                starts_name = any(word == a.split()[0] or a.startswith(word) and len(word) >= 3
                                  for al in spec["names"].values() for a in al)
                if word not in ARTICLE_EXCEPTIONS and not starts_name:
                    return mk("article_A", None, lone)
        label = after.split("\n", 1)[0][:LABEL_CHARS]
        label_named = _named_in(label, spec)
        label_letters = {next(g for g in mm.groups() if g) for mm in _STATED.finditer(label)} & set(letters)
        if (set(label_named) - {pick}) or (label_letters - {L}):   # any OTHER option named on the label line
            return mk("contradictory", None, None)
        return mk("letter+name" if label_named else "clean", pick, pick)

    if spec.get("json_pass"):
        j = _JSON_PASS.findall(s.lower())
        if j:
            return mk("json_verdict", None, "admit" if j[-1] == "true" else "refuse")

    stated = [next(g for g in mm.groups() if g) for mm in _STATED.finditer(s)]
    last = s.rstrip().split("\n")[-1]
    fm = _FINAL_LINE.match(last)
    if fm:
        stated.append(fm.group(1))
    stated = [x for x in stated if x in letters]
    if len(set(stated)) == 1:
        return mk("late_letter", None, opts[letters.index(stated[0])])
    if len(set(stated)) > 1:
        return mk("contradictory", None, None)
    if len(named_all) == 1:
        return mk("name_only", None, named_all[0])
    if len(named_all) > 1:
        return mk("contradictory", None, None)
    if any(ch in letters for ch in s):
        return mk("word_letter", None, None)
    return mk("unparsed", None, None)


def json_verdict_native(text: str) -> str | None:
    """The doorman's native contract (set id doorman36-json): the first top-level JSON object carrying a
    boolean `pass`; true -> admit, false -> refuse; anything else -> None (unparsed, scored wrong)."""
    t = (text or "").strip()
    if t.startswith("```"):
        t = re.sub(r"^```[a-zA-Z]*\s*|\s*```$", "", t)
    for start in (i for i, ch in enumerate(t) if ch == "{"):
        depth = 0
        for j in range(start, len(t)):
            if t[j] == "{":
                depth += 1
            elif t[j] == "}":
                depth -= 1
                if depth == 0:
                    try:
                        obj = json.loads(t[start:j + 1])
                    except json.JSONDecodeError:
                        break
                    if isinstance(obj, dict) and isinstance(obj.get("pass"), bool):
                        return "admit" if obj["pass"] else "refuse"
                    break
    return None


# ---------------------------------------------------------------------------------------------- selftest
#: (reply, set, stop_reason, expected class, expected strict choice, expected lenient choice, expected v1)
#: Every case the plan, the recon, the panel and the lanes found (PREREG-S0.md section 3 lists the source of each).
SELFTEST = [
    ("[F] socrates", "c", "stop", "letter+name", "socrates", "socrates", "socrates"),
    ("[F]", "c", "stop", "clean", "socrates", "socrates", "socrates"),
    ("\n[A] darwin: Charles", "c", "length", "letter+name", "darwin", "darwin", "darwin"),
    ("E", "c", "stop", "clean", "sagan", "sagan", "sagan"),
    ("B.", "c", "stop", "clean", "einstein", "einstein", "einstein"),
    ("**B**", "c", "stop", "clean", "einstein", "einstein", "einstein"),
    ("B\n\n{\"pass\": false}", "c", "stop", "clean", "einstein", "einstein", "einstein"),
    ("[C] hypatia: Hypat", "c", "length", "letter+name", "hypatia", "hypatia", "hypatia"),
    ("B) Albert Einstein", "c", "stop", "letter+name", "einstein", "einstein", "einstein"),
    ("[D] Ibn Sina", "c", "stop", "letter+name", "ibn_sina", "ibn_sina", "ibn_sina"),
    ("[E] sagan: Carl Sagan", "c", "stop", "letter+name", "sagan", "sagan", "sagan"),
    ("C. Socrates", "c", "stop", "contradictory", None, None, "hypatia"),
    ("[A] Darwin, not Socrates", "c", "stop", "contradictory", None, None, "darwin"),   # CHECK-PREREG MF-9 (a)
    ("A darwin", "c", "stop", "letter+name", "darwin", "darwin", "darwin"),
    ("A\nThe line is Darwin's", "c", "stop", "clean", "darwin", "darwin", "darwin"),
    ("A is the answer", "c", "stop", "clean", "darwin", "darwin", "darwin"),
    ("A) the line is about evolution", "c", "stop", "clean", "darwin", "darwin", "darwin"),   # CHECK-PREREG MF-9 (b)
    ("A line like this is Darwin", "c", "stop", "article_A", None, "darwin", "darwin"),
    ("A line like this…", "c", "length", "article_A", None, None, "darwin"),
    ("[A]\n\nThe passage references Socrates, not Darwin.", "c", "stop", "clean", "darwin", "darwin", "darwin"),
    ("**[D] ibn_sina: Ibn Sina**\n\nThe speaker addresses Socrates directly.", "c", "stop", "letter+name", "ibn_sina", "ibn_sina", "ibn_sina"),
    ("Einstein", "c", "stop", "name_only", None, "einstein", "sagan"),
    ("Carl Sagan", "c", "stop", "name_only", None, "sagan", "hypatia"),
    ("Darwin said it.", "c", "stop", "name_only", None, "darwin", "ibn_sina"),
    ("The speaker addresses Hypatia in the second", "c", "length", "name_only", None, "hypatia", "darwin"),
    ("Based on the style, [D] Ibn Sina", "c", "stop", "late_letter", None, "ibn_sina", "einstein"),
    ("Answer: B", "c", "stop", "late_letter", None, "einstein", "darwin"),
    ("\"Interconnected web of life\" is a Darwin motif.\n\nA", "c", "stop", "late_letter", None, "darwin", "ibn_sina"),
    ("The line addresses both Socrates and Hypatia, so it is neither of them.", "c", "stop", "contradictory", None, None, "darwin"),
    ("[B] or [C]", "c", "stop", "contradictory", None, None, "einstein"),
    ("We need answer letter only. Need infer", "c", "length", "unparsed", None, None, "darwin"),
    ("Certainly. The Down House reference settles it.", "c", "stop", "word_letter", None, None, "hypatia"),
    ("", "c", "length", "budget_exhausted", None, None, "darwin"),
    ("", "c", "stop", "unparsed", None, None, "darwin"),
    ("{\"pass\": false}", "c", "stop", "unparsed", None, None, "darwin"),
    ("b", "c", "stop", "unparsed", None, None, "darwin"),
    ("[A]", "a", "stop", "clean", "yes", "yes", "yes"),
    ("Based on the article, yes.", "a", "stop", "name_only", None, "yes", "no"),
    ("B", "d48", "stop", "clean", "refuse", "refuse", "refuse"),
    ("A", "d48", "stop", "clean", "admit", "admit", "admit"),
    ("[A] admit", "d48", "stop", "letter+name", "admit", "admit", "admit"),
    (" A\n\nExplanation: The line is fine.", "d48", "stop", "clean", "admit", "admit", "admit"),
    ("{\"pass\": true}", "d48", "stop", "json_verdict", None, "admit", "admit"),
    ("```json\n{\"pass\": false, \"class\": \"harassment\"}\n```", "d48", "stop", "json_verdict", None, "refuse", "admit"),
    ("I'm going to step out of the doorman role for this one.", "d48", "stop", "unparsed", None, None, "admit"),
    ("not grounded", "j", "stop", "name_only", None, "not_grounded", "grounded"),
    ("[C] off_page", "j", "stop", "letter+name", "off_page", "off_page", "off_page"),
    ("USER: [B]", "c", "stop", "late_letter", None, "einstein", "sagan"),
]


def run_selftest(as_table: bool = False) -> int:
    bad = []
    rows = []
    for text, sid, stop, cls, strict, lenient, v1 in SELFTEST:
        p = classify(text, sid, stop)
        ok = (p.format_class, p.choice_strict, p.choice_lenient, p.choice_v1) == (cls, strict, lenient, v1)
        if not ok:
            bad.append((text, sid, stop, (p.format_class, p.choice_strict, p.choice_lenient, p.choice_v1),
                        (cls, strict, lenient, v1)))
        rows.append((text, sid, stop, p))
    if as_table:
        print("| reply (verbatim, `⏎` = newline) | set | stop | class | strict | lenient | v1 (run.py) | template_suspect |")
        print("|---|---|---|---|---|---|---|---|")
        for text, sid, stop, p in rows:
            shown = text.replace("\n", "⏎").replace("|", "\\|")
            print("| `%s` | %s | %s | `%s` | %s | %s | %s | %s |" % (
                shown if shown else "(empty)", sid, stop, p.format_class,
                p.choice_strict or "—", p.choice_lenient or "—", p.choice_v1, "yes" if p.template_suspect else "no"))
    if bad:
        for b in bad:
            print("SELFTEST FAIL: %r set=%s stop=%s got=%s want=%s" % b, file=sys.stderr)
        return 1
    if not as_table:
        print("parse_v2 selftest: %d cases OK" % len(SELFTEST))
    return 0


def census(path: str, set_id: str, stop_field: str) -> dict:
    """Re-classify one lane row file. Refused rows (run.py's `refused`) stay out, as run.py keeps them."""
    counts: dict[str, int] = {}
    n = strict = lenient = v1 = 0
    for line in open(path, encoding="utf-8"):
        line = line.strip()
        if not line:
            continue
        r = json.loads(line)
        if r.get("refused"):
            continue
        p = classify(r.get("generated") or "", set_id, r.get(stop_field))
        counts[p.format_class] = counts.get(p.format_class, 0) + 1
        n += 1
        lab = r.get("label")
        strict += int(p.choice_strict is not None and p.choice_strict == lab)
        lenient += int(p.choice_lenient is not None and p.choice_lenient == lab)
        v1 += int(p.choice_v1 == lab)
    return {"file": path, "set": set_id, "n": n, "classes": dict(sorted(counts.items())),
            "correct_strict": strict, "correct_lenient": lenient, "correct_v1": v1,
            "parser_version": PARSER_VERSION, "parser_sha256": file_sha256()}


def file_sha256() -> str:
    with open(os.path.abspath(__file__), "rb") as fh:
        return hashlib.sha256(fh.read()).hexdigest()


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--selftest", action="store_true")
    ap.add_argument("--table", action="store_true", help="with --selftest: print the markdown table")
    ap.add_argument("--census", metavar="ROWS.jsonl")
    ap.add_argument("--set", default="c")
    ap.add_argument("--stop-field", default="done_reason")
    ap.add_argument("--sha", action="store_true", help="print this file's sha256")
    a = ap.parse_args()
    if a.sha:
        print(file_sha256())
        return 0
    if a.selftest:
        return run_selftest(a.table)
    if a.census:
        print(json.dumps(census(a.census, a.set, a.stop_field), indent=1))
        return 0
    ap.print_help()
    return 2


if __name__ == "__main__":
    sys.exit(main())