src/graphtrace/bench.py

"""Benchmark generation, human review and freezing.

Candidate questions are generated deterministically from the provenance graph with
fixed templates. Gold answers are *defined by source chunks*: every gold edge carries a
verbatim quote, and items are only released after a human checks the question, the
gold answer and the cited evidence against the source text (``graphtrace review``).
Generated items start as ``unreviewed``; ``freeze`` exports accepted/edited items only.

Threat to validity (see docs/PROTOCOL.md): questions are derived from the same graph
that graph-based methods retrieve from. Mitigations: gold is checked against source
text by a human, text-RAG sees the same source chunks, and extraction errors that slip
into gold are caught at review time (rejected items are reported).
"""

from __future__ import annotations

import hashlib
import json
import random
import re
from collections import defaultdict
from datetime import UTC, datetime
from pathlib import Path

from pydantic import BaseModel

from .canonical import mention_pattern, normalise_text
from .graph import KG, value_groups
from .llm import LLMClient, LLMError, call_structured
from .prompts import PARAPHRASE_SYSTEM, ParaphraseOut
from .schemas import BenchItem, Chunk, Edge, PathStep, Review


class BenchConfig(BaseModel):
    seed: int = 13
    targets: dict[str, int] = {"single_hop": 70, "multi_hop": 70, "false_premise": 40, "conflict_or_insufficient": 40}
    max_answer_set: int = 5
    min_edge_confidence: float | None = None
    popular_min_papers: int = 2  # an entity used for false premises must appear in >= this many papers


def _iid(*parts: object) -> str:
    return hashlib.sha1("|".join(map(str, parts)).encode("utf-8")).hexdigest()[:10]


def _step(e: Edge) -> PathStep:
    return PathStep(edge_id=e.edge_id, subject=e.subject, relation=e.relation, object=e.object)


class Generator:
    def __init__(self, kg: KG, chunks_by_doc: dict[str, list[Chunk]], cfg: BenchConfig):
        # Gold is computed on the same confidence-filtered graph that methods retrieve from
        # (KG.filtered), so both always see the same edge set.
        self.kg = kg.filtered(cfg.min_edge_confidence)
        self.cfg = cfg
        self.chunks_by_doc = chunks_by_doc
        self.chunk_text = {c.chunk_id: normalise_text(c.text) for cs in chunks_by_doc.values() for c in cs}
        self.doc_text = {d: " ".join(self.chunk_text[c.chunk_id] for c in cs) for d, cs in chunks_by_doc.items()}
        self.papers = sorted(n.node_id for n in self.kg.nodes.values() if n.type == "Paper")
        self.chunk_doc = {c.chunk_id: c.doc_id for cs in chunks_by_doc.values() for c in cs}

    # ---- helpers ----

    def objs(self, s: str, rel: str) -> dict[str, list[Edge]]:
        out: dict[str, list[Edge]] = defaultdict(list)
        for e in self.kg.out_edges(s, rel):
            out[e.object].append(e)
        return out

    def subs(self, o: str, rel: str) -> dict[str, list[Edge]]:
        out: dict[str, list[Edge]] = defaultdict(list)
        for e in self.kg.in_edges(o, rel):
            out[e.subject].append(e)
        return out

    def paper_datasets(self, pid: str) -> dict[str, list[Edge]]:
        """Datasets a paper evaluates on, directly or through a method it proposes.

        Assumption (stated in item notes): a paper's evaluation datasets and its proposed
        method's evaluation datasets are interchangeable. Method-derived entries carry the
        method's evaluates_on edge first and the proposes edge as provenance.
        """
        out: dict[str, list[Edge]] = defaultdict(list)
        for d, es in self.objs(pid, "evaluates_on").items():
            out[d].extend(es)
        for m, pe in self.objs(pid, "proposes").items():
            for d, es in self.objs(m, "evaluates_on").items():
                out[d].extend(es + pe)
        return out

    def forms(self, node_id: str) -> list[re.Pattern]:
        node = self.kg.nodes[node_id]
        pats = (mention_pattern(f) for f in {node.name, *node.aliases, node_id.split(":", 1)[1]})
        return [p for p in pats if p is not None]

    def title(self, pid: str) -> str:
        return self.kg.name(pid)

    def mentioned_in_paper(self, node_id: str, doc_id: str) -> bool:
        text = self.doc_text.get(doc_id, "")
        return any(p.search(text) for p in self.forms(node_id))

    def papers_of(self, node_id: str) -> set[str]:
        return {self.kg.nodes[e.subject].doc_id for e in self.kg.in_edges(node_id) if self.kg.nodes[e.subject].type == "Paper"}

    def item(
        self,
        category: str,
        template_id: str,
        hops: int,
        question: str,
        decision: str,
        answer_ids: list[str],
        edges: list[Edge],
        notes: str,
        answer_names: list[str] | None = None,
        extra_chunks: list[str] | None = None,
    ) -> BenchItem:
        names = answer_names if answer_names is not None else [self.kg.name(a) for a in answer_ids]
        chunks = sorted({e.source_chunk_id for e in edges} | set(extra_chunks or []))
        return BenchItem(
            item_id=f"{category[:2]}_{template_id}_{_iid(template_id, question)}",
            category=category,  # type: ignore[arg-type]
            difficulty=hops,
            template_id=template_id,
            question=question,
            question_template=question,
            gold_decision=decision,  # type: ignore[arg-type]
            gold_answer=names,
            gold_answer_ids=answer_ids,
            required_path=[_step(e) for e in edges],
            gold_chunk_ids=chunks,
            annotation_notes=notes,
        )

    @staticmethod
    def _first(es: list[Edge]) -> Edge:
        return min(es, key=lambda e: e.edge_id)

    def _set_item(self, cat, tid, hops, q, answer: dict[str, list[Edge]], extra_edges: list[Edge], notes) -> BenchItem | None:
        if not (1 <= len(answer) <= self.cfg.max_answer_set):
            return None
        ids = sorted(answer)
        edges = extra_edges + [self._first(answer[a]) for a in ids]
        prov = sorted({e.source_chunk_id for a in ids for e in answer[a]})
        return self.item(cat, tid, hops, q, "answer", ids, edges, notes, extra_chunks=prov)

    # ---- single hop ----

    def single_hop(self) -> list[BenchItem]:
        items = []
        spec = [
            ("S1", "evaluates_on", 'Which datasets does the paper "{t}" evaluate on?'),
            ("S2", "proposes", 'Which method does the paper "{t}" propose?'),
            ("S3", "addresses", 'Which task does the paper "{t}" address?'),
            ("S4", "reports_metric", 'Which evaluation metrics does the paper "{t}" report?'),
        ]
        for pid in self.papers:
            for tid, rel, tpl in spec:
                it = self._set_item(
                    "single_hop", tid, 1, tpl.format(t=self.title(pid)), self.objs(pid, rel), [], f"Gold = all '{rel}' objects of the paper in the graph."
                )
                if it:
                    items.append(it)
        # S5: unique result value lookup (same surface form as the conflict template C1)
        groups = defaultdict(list)
        for r in self.kg.results.values():
            groups[(r.method, r.dataset, r.metric)].append(r)
        for (m, d, k), rs in sorted(groups.items()):
            if all(r.value_num is not None for r in rs) and len(value_groups([r.value for r in rs])) == 1:
                q = f"What {self.kg.name(k)} does {self.kg.name(m)} achieve on {self.kg.name(d)}?"
                items.append(
                    self.item(
                        "single_hop",
                        "S5",
                        1,
                        q,
                        "answer",
                        [rs[0].value],
                        [],
                        "Gold = the reported value; all result records agree up to scale and rounding.",
                        answer_names=[rs[0].value],
                        extra_chunks=[r.source_chunk_id for r in rs],
                    ).model_copy(update={"required_path": [PathStep(edge_id=r.result_id, subject=m, relation=f"result:{k}", object=d) for r in rs]})
                )
        return items

    # ---- multi hop ----

    def multi_hop(self) -> list[BenchItem]:
        items = []
        pd = {p: self.paper_datasets(p) for p in self.papers}
        for pid in self.papers:
            t = self.title(pid)
            proposed = self.objs(pid, "proposes")
            # M1: P -proposes-> M -evaluates_on-> D
            if len(proposed) == 1:
                m, pe = next(iter(proposed.items()))
                it = self._set_item(
                    "multi_hop",
                    "M1",
                    2,
                    f'On which datasets is the method proposed in "{t}" evaluated?',
                    pd[pid],
                    [self._first(pe)],
                    "Path: paper proposes method; method evaluates_on dataset. Gold = union of the method's evaluates_on "
                    "datasets and the proposing paper's evaluates_on datasets (assumption: the paper's experiments evaluate "
                    "the method it proposes). REVIEWER: remove datasets used only for other purposes.",
                )
                if it:
                    items.append(it)
                # M4: P -proposes-> M -compares_with-> M2 <-proposes- P2
                ans: dict[str, list[Edge]] = defaultdict(list)
                hop2: list[Edge] = []
                for m2, ce in self.objs(m, "compares_with").items():
                    for p2, pe2 in self.subs(m2, "proposes").items():
                        if p2 != pid:
                            ans[p2].extend(pe2)
                            hop2.append(self._first(ce))
                it = self._set_item(
                    "multi_hop",
                    "M4",
                    3,
                    f'Which paper in the corpus proposes a method that the method proposed in "{t}" is compared with?',
                    ans,
                    [self._first(pe)] + sorted({e.edge_id: e for e in hop2}.values(), key=lambda e: e.edge_id),
                    "Path: paper proposes M; M compares_with M2; other paper proposes M2.",
                )
                if it:
                    items.append(it)
            # M2: P -evaluates_on-> D <-evaluates_on- P2
            ans2: dict[str, list[Edge]] = defaultdict(list)
            bridge: dict[str, Edge] = {}
            for d, de in pd[pid].items():
                for p2 in self.papers:
                    if p2 != pid and d in pd[p2]:
                        ans2[p2].extend(pd[p2][d])
                        bridge[d] = self._first(de)
            it = self._set_item(
                "multi_hop",
                "M2",
                2,
                f'Which other papers in the corpus evaluate on at least one of the same datasets as "{t}"?',
                ans2,
                sorted(bridge.values(), key=lambda e: e.edge_id),
                "Path: paper evaluates_on D; another paper evaluates_on D. A paper counts as evaluating on D directly or through a method it proposes.",
            )
            if it:
                items.append(it.model_copy(update={"gold_answer": [self.title(a) for a in it.gold_answer_ids]}))
        # M3: M <-uses- P -evaluates_on-> D  (aggregated over papers using M)
        for m in sorted(n.node_id for n in self.kg.nodes.values() if n.type == "Method"):
            users: dict[str, list[Edge]] = defaultdict(list)
            for rel in ("uses", "proposes"):
                for p, es in self.subs(m, rel).items():
                    if self.kg.nodes[p].type == "Paper":
                        users[p].extend(es)
            if len(users) < 2:
                continue
            ans3: dict[str, list[Edge]] = defaultdict(list)
            hop1 = []
            for p, ue in users.items():
                ds = pd[p]
                if ds:
                    hop1.append(self._first(ue))
                for d, de in ds.items():
                    ans3[d].extend(de)
            it = self._set_item(
                "multi_hop",
                "M3",
                2,
                f"Which datasets are evaluated on by papers in the corpus that use {self.kg.name(m)}?",
                ans3,
                sorted(hop1, key=lambda e: e.edge_id),
                "Path: paper uses or proposes M; that paper evaluates_on D, directly or through its proposed method (union over papers).",
            )
            if it:
                items.append(it)
        return items

    # ---- false premise ----

    def false_premise(self) -> list[BenchItem]:
        items = []
        popular_d = [n.node_id for n in self.kg.nodes.values() if n.type == "Dataset" and len(self.papers_of(n.node_id)) >= self.cfg.popular_min_papers]
        popular_m = [n.node_id for n in self.kg.nodes.values() if n.type == "Method" and len(self.papers_of(n.node_id)) >= self.cfg.popular_min_papers]
        rng = random.Random(self.cfg.seed)
        for pid in self.papers:
            doc = self.kg.nodes[pid].doc_id
            t = self.title(pid)
            proposed = self.objs(pid, "proposes")
            if len(proposed) == 1:
                _, pe = next(iter(proposed.items()))
                cands = [d for d in sorted(popular_d) if d not in self.paper_datasets(pid) and not self.mentioned_in_paper(d, doc)]
                if cands:
                    d = rng.choice(cands)
                    items.append(
                        self.item(
                            "false_premise",
                            "F1",
                            2,
                            f'What results does the method proposed in "{t}" achieve on {self.kg.name(d)}?',
                            "abstain",
                            [],
                            [self._first(pe)],
                            f"False premise: no evaluates_on edge from the paper or its method to {self.kg.name(d)}, and the dataset is not mentioned anywhere in the paper's text.",
                        )
                    )
            evals = self.objs(pid, "evaluates_on")
            if evals:
                used = set(self.objs(pid, "uses")) | set(proposed)
                cands = [m for m in sorted(popular_m) if m not in used and not self.mentioned_in_paper(m, doc)]
                if cands:
                    m2 = rng.choice(cands)
                    items.append(
                        self.item(
                            "false_premise",
                            "F2",
                            2,
                            f'Which datasets did the paper "{t}" use to evaluate {self.kg.name(m2)}?',
                            "abstain",
                            [],
                            [self._first(min(evals.values(), key=lambda es: es[0].edge_id))],
                            f"False premise: the paper neither proposes nor uses {self.kg.name(m2)} and never mentions it.",
                        )
                    )
        return items

    # ---- conflict / insufficient ----

    def conflict_or_insufficient(self) -> list[BenchItem]:
        items = []
        groups = defaultdict(list)
        for r in self.kg.results.values():
            groups[(r.method, r.dataset, r.metric)].append(r)
        for (m, d, k), rs in sorted(groups.items()):
            numeric = [r for r in rs if r.value_num is not None]
            vgroups = value_groups(sorted({r.value for r in numeric}))
            chunks = {r.source_chunk_id for r in numeric}
            if len(vgroups) >= 2 and len(chunks) >= 2:
                q = f"What {self.kg.name(k)} does {self.kg.name(m)} achieve on {self.kg.name(d)}?"
                values = sorted({r.value for r in numeric})
                docs = {self.chunk_doc.get(r.source_chunk_id) for r in numeric}
                scope = (
                    "Sources disagree across papers: the values come from different papers (for example a paper's own "
                    "number and another paper's reproduction or citation). This counts as a disagreement. "
                    if len(docs) > 1
                    else "The values come from different passages of the same paper. "
                )
                items.append(
                    self.item(
                        "conflict_or_insufficient",
                        "C1",
                        1,
                        q,
                        "flag_conflict",
                        values,
                        [],
                        "Conflict: values that differ after scale normalisation and rounding to the reported precision, for "
                        "the same (method, dataset, metric). " + scope + "REVIEWER: reject if the values refer to different "
                        "settings (data split, model size, metric variant).",
                        answer_names=values,
                        extra_chunks=sorted(chunks),
                    ).model_copy(update={"required_path": [PathStep(edge_id=r.result_id, subject=m, relation=f"result:{k}", object=d) for r in numeric]})
                )
        # C2: method evaluated on dataset, but no reported value anywhere in the corpus
        for e in sorted(self.kg.edges.values(), key=lambda e: e.edge_id):
            if e.relation != "evaluates_on" or self.kg.nodes[e.subject].type != "Method":
                continue
            m, d = e.subject, e.object
            if self.kg.results_for(method=m, dataset=d):
                continue
            mf, df = self.forms(m), self.forms(d)
            # A passage naming both (any alias, same normalisation on both sides) next to any
            # number, integer or decimal, might report a value: not safely "insufficient".
            if any(any(p.search(t) for p in mf) and any(p.search(t) for p in df) and re.search(r"\d", t) for t in self.chunk_text.values()):
                continue
            doc = self.chunk_doc.get(e.source_chunk_id)
            paper_metrics = self.objs(f"paper:{doc}", "reports_metric") if doc else {}
            if not paper_metrics:
                continue  # ask only about a metric the source paper actually reports
            k = min(paper_metrics)
            items.append(
                self.item(
                    "conflict_or_insufficient",
                    "C2",
                    1,
                    f"What {self.kg.name(k)} does {self.kg.name(m)} achieve on {self.kg.name(d)}?",
                    "abstain",
                    [],
                    [e],
                    "Insufficient evidence: the method is evaluated on the dataset, but no passage in the corpus states a value for this pair. "
                    "REVIEWER: confirm no table/figure caption gives the number.",
                )
            )
        return items

    # ---- sampling ----

    def generate(self) -> tuple[list[BenchItem], dict]:
        pools = {
            "single_hop": self.single_hop(),
            "multi_hop": self.multi_hop(),
            "false_premise": self.false_premise(),
            "conflict_or_insufficient": self.conflict_or_insufficient(),
        }
        rng = random.Random(self.cfg.seed)
        chosen: list[BenchItem] = []
        stats: dict = {"available": {}, "selected": {}, "by_template": {}}
        for cat, items in pools.items():
            uniq = {it.item_id: it for it in items}
            by_t: dict[str, list[BenchItem]] = defaultdict(list)
            for it in sorted(uniq.values(), key=lambda i: i.item_id):
                by_t[it.template_id].append(it)
            for lst in by_t.values():
                rng.shuffle(lst)
            stats["available"][cat] = {t: len(v) for t, v in sorted(by_t.items())}
            # Round-robin across templates so no single template dominates a category.
            picked: list[BenchItem] = []
            target = self.cfg.targets.get(cat, 0)
            order = sorted(by_t)
            while len(picked) < target and any(by_t[t] for t in order):
                for t in order:
                    if by_t[t] and len(picked) < target:
                        picked.append(by_t[t].pop())
            stats["selected"][cat] = len(picked)
            for it in picked:
                stats["by_template"][it.template_id] = stats["by_template"].get(it.template_id, 0) + 1
            chosen.extend(picked)
        return sorted(chosen, key=lambda i: i.item_id), stats


# ---- paraphrase (optional, LLM) --------------------------------------------------


def protected_spans(item: BenchItem) -> list[str]:
    spans = re.findall(r'"([^"]+)"', item.question_template)
    return spans


def paraphrase(client: LLMClient, item: BenchItem, required_names: list[str]) -> tuple[str, str | None]:
    """Return (question, rejection reason). Falls back to the template on any problem."""
    try:
        out = call_structured(client, None, "paraphrase", PARAPHRASE_SYSTEM, item.question_template, ParaphraseOut, max_tokens=512)
    except LLMError as e:
        return item.question_template, f"llm_error: {e}"
    q = out.question.strip()
    missing = [n for n in protected_spans(item) + required_names if n and n not in q]
    if missing:
        return item.question_template, f"dropped entity names: {missing}"
    return q, None


# ---- review & freeze ----------------------------------------------------------------


def read_items(path: str | Path) -> list[BenchItem]:
    p = Path(path)
    return [BenchItem(**json.loads(x)) for x in p.read_text(encoding="utf-8").splitlines() if x.strip()]


def write_items(path: str | Path, items: list[BenchItem]) -> None:
    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    with p.open("w", encoding="utf-8") as f:
        for it in items:
            f.write(it.model_dump_json() + "\n")


def read_reviews(path: str | Path) -> list[Review]:
    p = Path(path)
    if not p.exists():
        return []
    return [Review(**json.loads(x)) for x in p.read_text(encoding="utf-8").splitlines() if x.strip()]


def append_review(path: str | Path, review: Review) -> None:
    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    with p.open("a", encoding="utf-8") as f:
        f.write(review.model_dump_json() + "\n")


def make_review(item_id: str, status: str, reviewer: str, notes: str = "", **edits) -> Review:
    return Review(item_id=item_id, status=status, reviewer=reviewer, notes=notes, timestamp=datetime.now(UTC).isoformat(), **edits)  # type: ignore[arg-type]


EDIT_NOTE = " [gold_answer edited in review: required_path and gold_chunk_ids were not regenerated]"


def apply_reviews(items: list[BenchItem], reviews: list[Review], linker=None) -> list[BenchItem]:
    """Apply every review in timestamp order (stable for equal timestamps).

    Status, reviewer and notes come from the latest review; edits accumulate, so a later
    "accepted" keeps earlier edits. An edited gold answer is re-linked to graph nodes with
    ``linker`` (names that do not link are kept as literal strings); the item notes that its
    required path and gold chunks were not regenerated.
    """
    by_item: dict[str, list[Review]] = defaultdict(list)
    for r in sorted(reviews, key=lambda r: r.timestamp):
        by_item[r.item_id].append(r)
    out = []
    for it in items:
        rs = by_item.get(it.item_id)
        if not rs:
            out.append(it)
            continue
        upd: dict = {}
        for r in rs:
            upd.update(review_status=r.status, reviewer=r.reviewer, review_notes=r.notes)
            if r.question is not None:
                upd["question"] = r.question
            if r.gold_decision is not None:
                upd["gold_decision"] = r.gold_decision
            if r.gold_answer is not None:
                upd["gold_answer"] = r.gold_answer
                ids = []
                for name in r.gold_answer:
                    lk = linker.link_name(name) if linker is not None else None
                    ids.append(lk.node_id if lk else name)
                upd["gold_answer_ids"] = ids
                if EDIT_NOTE not in it.annotation_notes:
                    upd["annotation_notes"] = it.annotation_notes + EDIT_NOTE
        out.append(it.model_copy(update=upd))
    return out


def freeze(items: list[BenchItem], reviews: list[Review], linker=None) -> tuple[list[BenchItem], dict]:
    reviewed = apply_reviews(items, reviews, linker)
    edited_ids = {r.item_id for r in reviews if r.status == "edited"}
    release = [it for it in reviewed if it.review_status in ("accepted", "edited")]
    by_cat: dict[str, int] = defaultdict(int)
    for it in release:
        by_cat[it.category] += 1
    manifest = {
        "frozen_utc": datetime.now(UTC).isoformat(),
        "candidates": len(items),
        "reviewed": sum(1 for it in reviewed if it.review_status != "unreviewed"),
        "released": len(release),
        "rejected": sum(1 for it in reviewed if it.review_status == "rejected"),
        "edited": sum(1 for it in release if it.item_id in edited_ids),
        "by_category": dict(by_cat),
        "reviewers": sorted({it.reviewer for it in release if it.reviewer}),
    }
    return release, manifest