core/generate/meaning_graph/projectors.py

"""Projectors — map a comprehended ``MeaningGraph`` into a reasoner's input shape.

The comprehension reader produces a neutral ``MeaningGraph``; a projector adapts
it to a specific independent-gold reasoner so the reader can be scored end-to-end
(prose -> MeaningGraph -> projection -> oracle -> answer vs gold). Projectors hold
NO decision logic — they only re-shape; the verdict is the independent oracle's.
"""

from __future__ import annotations

from typing import Any

from generate.meaning_graph.reader import Comprehension


def to_set_membership(comp: Comprehension) -> tuple[dict[str, Any], dict[str, Any]] | None:
    """Project into ``(structure, query)`` for ``evals.set_membership.oracle``.

    Returns ``None`` when the comprehension does not carry exactly one membership
    question (nothing to ask the oracle) — the caller treats that as a refusal.
    """
    graph = comp.meaning_graph
    individuals = sorted({e.entity_id for e in graph.entities if e.kind == "individual"})
    classes = sorted({e.entity_id for e in graph.entities if e.kind == "class"})

    member_facts = [
        (r.arguments[0], r.arguments[1])
        for r in graph.relations
        if r.predicate == "member" and not r.negated
    ]
    subset_facts = [
        (r.arguments[0], r.arguments[1])
        for r in graph.relations
        if r.predicate == "subset" and not r.negated
    ]

    member_queries = [q for q in comp.queries if q.predicate == "member" and not q.negated]
    subset_queries = [q for q in comp.queries if q.predicate == "subset" and not q.negated]
    if len(member_queries) + len(subset_queries) != 1:
        return None

    sets = [
        {"id": cid, "members": sorted({ind for ind, cls in member_facts if cls == cid})}
        for cid in classes
    ]
    structure = {
        "elements": individuals,
        "sets": sets,
        "subsets": [{"subset": a, "superset": b} for a, b in subset_facts],
    }
    if member_queries:
        q = member_queries[0]
        query = {"kind": "member", "element": q.arguments[0], "set": q.arguments[1]}
    else:
        q = subset_queries[0]
        query = {"kind": "subset", "subset": q.arguments[0], "superset": q.arguments[1]}
    return structure, query


#: Categorical predicate -> Aristotelian form for the syllogism oracle.
_PRED_FORM = {"subset": "A", "disjoint": "E", "intersects": "I", "some_not": "O"}

#: Finite-model domain size for syllogism validity (matches the gold lane).
_SYLLOGISM_DOMAIN_SIZE = 3


def to_syllogism(comp: Comprehension) -> tuple[dict[str, Any], dict[str, Any]] | None:
    """Project into ``(structure, query)`` for ``evals.syllogism.oracle``.

    Premises are the categorical relations (subset/disjoint/intersects/some_not);
    the single categorical query is the conclusion. Returns ``None`` when the
    comprehension is not exactly one categorical conclusion over >=1 premise — the
    caller treats that as a refusal (nothing honestly askable of this oracle).
    """
    graph = comp.meaning_graph
    premises = [
        {"form": _PRED_FORM[r.predicate], "subject": r.arguments[0], "predicate": r.arguments[1]}
        for r in graph.relations
        if r.predicate in _PRED_FORM and not r.negated
    ]
    conclusions = [q for q in comp.queries if q.predicate in _PRED_FORM and not q.negated]
    if not premises or len(comp.queries) != 1 or len(conclusions) != 1:
        return None

    c = conclusions[0]
    terms = sorted({e.entity_id for e in graph.entities if e.kind == "class"})
    if len(terms) < 2:
        return None
    structure = {
        "terms": terms,
        "domain_size": _SYLLOGISM_DOMAIN_SIZE,
        "premises": premises,
    }
    query = {
        "kind": "validity",
        "conclusion": {
            "form": _PRED_FORM[c.predicate],
            "subject": c.arguments[0],
            "predicate": c.arguments[1],
        },
    }
    return structure, query


def to_total_ordering(comp: Comprehension) -> tuple[dict[str, Any], dict[str, Any]] | None:
    """Project into ``(structure, query)`` for ``evals.total_ordering.oracle``.

    Facts are ``less(lo, hi)`` relations; the single query is a sort or compare.
    Returns ``None`` when the comprehension does not carry exactly one ordering
    query — the caller treats that as a refusal.
    """
    graph = comp.meaning_graph
    less_facts = [
        (r.arguments[0], r.arguments[1])
        for r in graph.relations
        if r.predicate == "less" and not r.negated
    ]
    order_queries = [q for q in comp.queries if q.predicate in ("sort", "compare")]
    if len(comp.queries) != 1 or len(order_queries) != 1:
        return None

    q = order_queries[0]
    item_set = {item for pair in less_facts for item in pair}
    if q.predicate == "compare":
        item_set.update(q.arguments)
    structure = {
        "items": sorted(item_set),
        "relations": [{"less": lo, "greater": hi} for lo, hi in less_facts],
    }
    if q.predicate == "sort":
        query = {"kind": "sort", "order": q.arguments[0]}
    else:
        query = {"kind": "compare", "left": q.arguments[0], "right": q.arguments[1]}
    return structure, query