arborist/bench/batteries/b_5t.py

"""5T battery — Transitivity / Transfer Learning / Triangulation /
Truthtables / Time.

Phase 1a (landed in ticket #000021): Transfer (legacy SQD-whitepaper
name). Stays intact for back-compat.

Phase 1b (ticket #000024, this file):
- Vocabulary aligned to Dav1DPrometheus: ``transfer-learning``,
  ``triangulation``, ``truthtables``, ``transitivity``, ``time``.
- ``transfer-v1`` digest stays pinned; new ``transfer-learning-v2``
  is the canonical name going forward.
- Five non-stub runners over text / claim_lattice / propositional /
  memory-root carriers.

No LLM-as-judge anywhere. All runners deterministic. Cross-modality
discipline: fixtures may carry future-modality hooks; v1 carriers
limited to the Phase-1 whitelist in ``base.PHASE_1_CARRIERS``.
"""

from __future__ import annotations

import json
import re
from itertools import permutations
from pathlib import Path
from typing import Iterable

from bench.batteries.base import (
    BatteryResult,
    TaskResult,
    fixture_digest,
    fixture_meta,
    iter_tasks,
    validate_carrier,
)
from bench.batteries.b_5s import _runtime_digest


def _carrier_check(task: dict) -> tuple[bool, str]:
    reason = validate_carrier(task)
    if reason is not None:
        return False, reason
    return True, ""


def _build_result(
    sub_battery: str,
    fixtures_path: Path,
    per_task: list[TaskResult],
    metrics: dict[str, float],
) -> BatteryResult:
    pass_count = sum(1 for t in per_task if t.passed)
    fail_count = len(per_task) - pass_count
    return BatteryResult(
        battery="5t",
        sub_battery=sub_battery,
        fixture_path=str(fixtures_path),
        fixture_digest=fixture_digest(fixtures_path),
        pass_count=pass_count,
        fail_count=fail_count,
        metrics=metrics,
        per_task=per_task,
        runtime_digest=_runtime_digest(),
    )


# ---------------------------------------------------------------------
# Phase 1a: Transfer (legacy)
# ---------------------------------------------------------------------


def run_transfer(fixtures_path: Path) -> BatteryResult:
    """Legacy SQD-whitepaper Transfer runner.

    Each task: π* + two surface variants → does canonicalize collapse
    them as expected_equivalent says? Kept untouched for Phase 1a
    digest stability.
    """
    from arborist.pi_star import get

    per_task: list[TaskResult] = []
    for task in iter_tasks(fixtures_path):
        task_id = task["id"]
        ok, reason = _carrier_check(task)
        if not ok:
            per_task.append(
                TaskResult(task_id=task_id, passed=False, detail={"reason": reason})
            )
            continue
        detail: dict = {}
        try:
            pi_star = get(task["pi_star"])
            ca = pi_star.canonicalize(task["input_a"].encode("utf-8"))
            cb = pi_star.canonicalize(task["input_b"].encode("utf-8"))
            equivalent = ca == cb
            ok = equivalent == bool(task["expected_equivalent"])
            if not ok:
                detail["reason"] = (
                    f"observed_equivalent={equivalent} "
                    f"expected={task['expected_equivalent']}"
                )
        except Exception as exc:  # noqa: BLE001
            detail["reason"] = f"{type(exc).__name__}: {exc}"
            ok = False
        per_task.append(TaskResult(task_id=task_id, passed=ok, detail=detail))

    total = len(per_task)
    rate = sum(1 for t in per_task if t.passed) / total if total else 0.0
    meta = fixture_meta(fixtures_path)
    return _build_result(
        meta.get("sub_battery", "transfer"),
        fixtures_path,
        per_task,
        {"transfer_pass_rate": rate},
    )


# ---------------------------------------------------------------------
# Phase 1b: Transfer Learning (Dav1DPrometheus canonical)
# ---------------------------------------------------------------------


def run_transfer_learning(fixtures_path: Path) -> BatteryResult:
    """Pattern transfer across tasks/domains.

    Each task: a source-task pattern + a target-task pattern + an
    expected_transfer flag. Runner asserts that the target task
    can be expressed under the same canonical structure as the
    source — for v1, this means the target's content tokens are a
    superset of the source's content tokens at the *pattern* level
    (not the example level).
    """
    per_task: list[TaskResult] = []
    for task in iter_tasks(fixtures_path):
        task_id = task["id"]
        ok, reason = _carrier_check(task)
        if not ok:
            per_task.append(
                TaskResult(task_id=task_id, passed=False, detail={"reason": reason})
            )
            continue
        try:
            src_pattern = task["source_task"]["pattern"]
            tgt_pattern = task["target_task"]["pattern"]
            expected = bool(task["target_task"].get("expected_transfer", True))
            structural_match = _patterns_share_structure(src_pattern, tgt_pattern)
            passed = structural_match == expected
            detail = {
                "expected_transfer": expected,
                "observed_transfer": structural_match,
            }
        except Exception as exc:  # noqa: BLE001
            passed = False
            detail = {"reason": f"{type(exc).__name__}: {exc}"}
        per_task.append(TaskResult(task_id=task_id, passed=passed, detail=detail))

    total = len(per_task)
    rate = sum(1 for t in per_task if t.passed) / total if total else 0.0
    meta = fixture_meta(fixtures_path)
    return _build_result(
        meta.get("sub_battery", "transfer-learning"),
        fixtures_path,
        per_task,
        {"transfer_learning_success_rate": rate},
    )


def _patterns_share_structure(a: str, b: str) -> bool:
    """Two patterns share structure iff they share ≥ 50% of content
    tokens AND have similar length (within 50% of each other)."""
    from bench.batteries.b_5s import _content_tokens

    a_tokens = _content_tokens(a)
    b_tokens = _content_tokens(b)
    if not a_tokens or not b_tokens:
        return False
    overlap = len(a_tokens & b_tokens) / max(len(a_tokens), len(b_tokens))
    len_ratio = min(len(a), len(b)) / max(len(a), len(b))
    return overlap >= 0.5 and len_ratio >= 0.5


# ---------------------------------------------------------------------
# Phase 1b: Triangulation
# ---------------------------------------------------------------------

# Independent verifier strategies that agree (or don't) on whether a
# claim is supported by evidence text. Each is deterministic.

def _strat_substring(claim: str, evidence: str) -> bool:
    return claim.strip().lower() in evidence.lower()


def _strat_token_subset(claim: str, evidence: str) -> bool:
    from bench.batteries.b_5s import _content_tokens

    a = _content_tokens(claim)
    b = _content_tokens(evidence)
    return bool(a) and a.issubset(b)


def _strat_token_overlap(claim: str, evidence: str, threshold: float = 0.6) -> bool:
    from bench.batteries.b_5s import _content_tokens

    a = _content_tokens(claim)
    b = _content_tokens(evidence)
    if not a:
        return False
    return len(a & b) / len(a) >= threshold


def _strat_entity_match(claim: str, evidence: str) -> bool:
    """Capitalized tokens (proper nouns) in claim must appear in evidence."""
    entities = {
        tok.strip(".,;:!?")
        for tok in claim.split()
        if tok and tok[0].isupper() and tok.lower() not in {"a", "an", "the", "i"}
    }
    if not entities:
        return False
    e_low = evidence.lower()
    return all(e.lower() in e_low for e in entities)


_TRIANGULATION_STRATEGIES = {
    "substring": _strat_substring,
    "token_subset": _strat_token_subset,
    "token_overlap": _strat_token_overlap,
    "entity_match": _strat_entity_match,
}


def run_triangulation(fixtures_path: Path) -> BatteryResult:
    """Independent-strategy agreement on (claim, evidence) pairs.

    Each task names a list of strategies + an ``expected_agreement_min``.
    Runner runs each strategy, computes pass-rate across them, and
    accepts the task iff agreement rate ≥ threshold.
    """
    per_task: list[TaskResult] = []
    for task in iter_tasks(fixtures_path):
        task_id = task["id"]
        ok, reason = _carrier_check(task)
        if not ok:
            per_task.append(
                TaskResult(task_id=task_id, passed=False, detail={"reason": reason})
            )
            continue
        strategies = task.get("strategies", list(_TRIANGULATION_STRATEGIES))
        claim = task["claim"]
        evidence = task["evidence"]
        results = {}
        verdicts: list[bool] = []
        for s in strategies:
            fn = _TRIANGULATION_STRATEGIES.get(s)
            if fn is None:
                continue
            v = fn(claim, evidence)
            results[s] = v
            verdicts.append(v)
        if not verdicts:
            per_task.append(
                TaskResult(task_id=task_id, passed=False,
                           detail={"reason": "no recognized strategies"})
            )
            continue
        agreement = sum(verdicts) / len(verdicts)
        threshold = float(task.get("expected_agreement_min", 0.75))
        expected = task.get("expected", "pass")
        # expected="pass" → claim should triangulate; agreement>=threshold.
        # expected="fail" → claim should NOT triangulate; agreement<threshold.
        if expected == "pass":
            passed = agreement >= threshold
        else:
            passed = agreement < threshold
        per_task.append(
            TaskResult(
                task_id=task_id,
                passed=passed,
                detail={
                    "strategy_results": results,
                    "agreement": agreement,
                    "threshold": threshold,
                    "expected": expected,
                },
            )
        )

    total = len(per_task)
    rate = sum(1 for t in per_task if t.passed) / total if total else 0.0
    meta = fixture_meta(fixtures_path)
    return _build_result(
        meta.get("sub_battery", "triangulation"),
        fixtures_path,
        per_task,
        {"triangulation_agreement_rate": rate},
    )


# ---------------------------------------------------------------------
# Phase 1b: Truthtables (N=2..4)
# ---------------------------------------------------------------------


def _eval_propositional(expr: str, env: dict[str, bool]) -> bool:
    """Deterministic propositional evaluator for AND/OR/NOT/XOR.

    Recognizes tokens A-Z (variable), AND, OR, NOT, XOR, IMPL, IFF,
    parentheses. No LLM. Operator precedence: NOT > AND > OR/XOR/IMPL/IFF
    (left-to-right within same precedence). Parens override.

    Implementation: tokenize → recursive-descent parse → evaluate.
    """
    tokens = _tokenize_propositional(expr)
    pos = [0]

    def expect(t: str) -> None:
        if pos[0] >= len(tokens) or tokens[pos[0]] != t:
            raise ValueError(f"expected {t!r} at {pos[0]}")
        pos[0] += 1

    def peek() -> str | None:
        if pos[0] >= len(tokens):
            return None
        return tokens[pos[0]]

    def parse_atom() -> bool:
        t = peek()
        if t is None:
            raise ValueError("unexpected end of expression")
        if t == "(":
            pos[0] += 1
            v = parse_or()
            expect(")")
            return v
        if t == "NOT":
            pos[0] += 1
            return not parse_atom()
        if t in ("TRUE", "FALSE"):
            pos[0] += 1
            return t == "TRUE"
        if len(t) == 1 and t.isupper() and t.isalpha():
            if t not in env:
                raise ValueError(f"variable {t!r} not in env")
            pos[0] += 1
            return env[t]
        raise ValueError(f"unexpected token {t!r}")

    def parse_and() -> bool:
        left = parse_atom()
        while peek() == "AND":
            pos[0] += 1
            right = parse_atom()
            left = left and right
        return left

    def parse_or() -> bool:
        left = parse_and()
        while peek() in ("OR", "XOR", "IMPL", "IFF"):
            op = peek()
            pos[0] += 1
            right = parse_and()
            if op == "OR":
                left = left or right
            elif op == "XOR":
                left = left != right
            elif op == "IMPL":
                left = (not left) or right
            else:  # IFF
                left = left == right
        return left

    val = parse_or()
    if pos[0] != len(tokens):
        raise ValueError(f"unconsumed tokens at {pos[0]}: {tokens[pos[0]:]!r}")
    return val


def _tokenize_propositional(expr: str) -> list[str]:
    """Split an expression into tokens, treating AND/OR/NOT/XOR/IMPL/IFF as keywords."""
    out: list[str] = []
    i = 0
    while i < len(expr):
        c = expr[i]
        if c.isspace():
            i += 1
            continue
        if c in "()":
            out.append(c)
            i += 1
            continue
        if c.isalpha():
            j = i
            while j < len(expr) and expr[j].isalpha():
                j += 1
            tok = expr[i:j]
            up = tok.upper()
            if up in ("AND", "OR", "NOT", "XOR", "IMPL", "IFF", "TRUE", "FALSE"):
                out.append(up)
            elif len(tok) == 1:
                out.append(tok.upper())
            else:
                raise ValueError(f"unrecognized token: {tok!r}")
            i = j
            continue
        raise ValueError(f"unexpected char: {c!r}")
    return out


def run_truthtables(fixtures_path: Path) -> BatteryResult:
    """Exhaustive propositional truth-table evaluation; N=2..4.

    Each task lists `variables` (≤4), an `expression`, and a `rows`
    array enumerating expected outputs. Runner evaluates the
    expression on every row and fails the task if any row mismatches.
    """
    per_task: list[TaskResult] = []
    coverage_total = 0
    coverage_correct = 0
    for task in iter_tasks(fixtures_path):
        task_id = task["id"]
        ok, reason = _carrier_check(task)
        if not ok:
            per_task.append(
                TaskResult(task_id=task_id, passed=False, detail={"reason": reason})
            )
            continue
        variables = task["variables"]
        if len(variables) > 4:
            per_task.append(
                TaskResult(task_id=task_id, passed=False,
                           detail={"reason": f"N={len(variables)} exceeds cap of 4"})
            )
            continue
        rows = task["rows"]
        coverage_total += len(rows)
        all_correct = True
        first_mismatch = None
        for row in rows:
            try:
                got = _eval_propositional(task["expression"], row["inputs"])
            except Exception as exc:  # noqa: BLE001
                all_correct = False
                first_mismatch = {
                    "inputs": row["inputs"],
                    "error": f"{type(exc).__name__}: {exc}",
                }
                break
            if got == bool(row["expected"]):
                coverage_correct += 1
            else:
                all_correct = False
                if first_mismatch is None:
                    first_mismatch = {
                        "inputs": row["inputs"],
                        "expected": row["expected"],
                        "got": got,
                    }
        detail: dict = {"row_count": len(rows)}
        if first_mismatch is not None:
            detail["first_mismatch"] = first_mismatch
        per_task.append(TaskResult(task_id=task_id, passed=all_correct, detail=detail))

    total = len(per_task)
    pass_rate = sum(1 for t in per_task if t.passed) / total if total else 0.0
    row_acc = coverage_correct / coverage_total if coverage_total else 0.0
    meta = fixture_meta(fixtures_path)
    return _build_result(
        meta.get("sub_battery", "truthtables"),
        fixtures_path,
        per_task,
        {
            "truth_table_coverage_rate": pass_rate,
            "row_accuracy_rate": row_acc,
        },
    )


# ---------------------------------------------------------------------
# Phase 1b: Transitivity (typed-relation whitelist)
# ---------------------------------------------------------------------

# Per ticket #000024 §4.3: only relations on this whitelist license
# transitive closure. "related_to" / "causes" without domain
# constraints are NOT transitive.
_TRANSITIVE_RELATIONS = frozenset({
    "implies",
    "subset_of",
    "ancestor_of",
    "before",
    "less_than",
})


def _walk_relation_path(
    edges: list[dict], start: str, end: str, relation: str
) -> bool:
    """BFS over edges of the given relation only. Returns True iff a
    path from start to end exists in the relation-restricted subgraph.
    """
    if relation not in _TRANSITIVE_RELATIONS:
        return False
    adj: dict[str, set[str]] = {}
    for e in edges:
        if e.get("relation") != relation:
            continue
        adj.setdefault(e["from"], set()).add(e["to"])
    if start == end:
        return True
    visited = {start}
    queue = [start]
    while queue:
        node = queue.pop(0)
        for nxt in adj.get(node, ()):
            if nxt == end:
                return True
            if nxt not in visited:
                visited.add(nxt)
                queue.append(nxt)
    return False


def run_transitivity(fixtures_path: Path) -> BatteryResult:
    """Typed-relation chain pass-rate.

    Each task: typed edges + a query (from, to, relation). Runner
    walks the edge subgraph for that relation and asserts whether a
    path exists. Mixed-relation chains never pass — runner only
    considers edges of the queried relation. Non-transitive
    relations always fail by construction.
    """
    per_task: list[TaskResult] = []
    for task in iter_tasks(fixtures_path):
        task_id = task["id"]
        ok, reason = _carrier_check(task)
        if not ok:
            per_task.append(
                TaskResult(task_id=task_id, passed=False, detail={"reason": reason})
            )
            continue
        try:
            q = task["query"]
            edges = task["edges"]
            path_exists = _walk_relation_path(edges, q["from"], q["to"], q["relation"])
            observed = "pass" if path_exists else "fail"
            expected = task.get("expected", "pass")
            passed = observed == expected
            detail = {
                "expected": expected,
                "observed": observed,
                "relation": q["relation"],
                "transitive": q["relation"] in _TRANSITIVE_RELATIONS,
            }
        except Exception as exc:  # noqa: BLE001
            passed = False
            detail = {"reason": f"{type(exc).__name__}: {exc}"}
        per_task.append(TaskResult(task_id=task_id, passed=passed, detail=detail))

    total = len(per_task)
    rate = sum(1 for t in per_task if t.passed) / total if total else 0.0
    meta = fixture_meta(fixtures_path)
    return _build_result(
        meta.get("sub_battery", "transitivity"),
        fixtures_path,
        per_task,
        {"full_chain_pass_rate": rate},
    )


# ---------------------------------------------------------------------
# Phase 1b: Time (memory-root snapshots)
# ---------------------------------------------------------------------


def run_time(fixtures_path: Path) -> BatteryResult:
    """Temporal-context preservation across snapshots.

    Each task: an ordered list of snapshots (each with a memory_root,
    facts, optional stale-marker) + an `expected` block stating which
    facts should preserve, which should mark stale, and what the
    final root is. Runner deterministically validates these
    properties from the fixture itself — no shard reads required for
    Phase 1b. Real shard-driven Time runs ship in Phase 1b.2 once
    arborist's memory_records is populated by a real workload.
    """
    per_task: list[TaskResult] = []
    for task in iter_tasks(fixtures_path):
        task_id = task["id"]
        ok, reason = _carrier_check(task)
        if not ok:
            per_task.append(
                TaskResult(task_id=task_id, passed=False, detail={"reason": reason})
            )
            continue
        try:
            snapshots = task["snapshots"]
            expected = task["expected"]
            # All facts seen across the whole chain.
            all_facts: set[str] = set()
            for s in snapshots:
                all_facts.update(s.get("facts", []))
                all_facts.update(s.get("stale_facts", []))
            # Preservation check: every fact named in expected.preserved_facts
            # must appear in some snapshot.
            preserved_ok = all(
                f in all_facts
                for f in expected.get("preserved_facts", [])
            )
            # Stale check: every fact in expected.marks_stale must appear
            # at least once and the FINAL snapshot's stale_facts must
            # include it.
            final_stale = set(snapshots[-1].get("stale_facts", []))
            stale_ok = all(
                f in final_stale
                for f in expected.get("marks_stale", [])
            )
            # Current root check.
            current_root_ok = (
                snapshots[-1].get("memory_root") == expected["current_root"]
            )
            passed = preserved_ok and stale_ok and current_root_ok
            detail = {
                "preserved_ok": preserved_ok,
                "stale_ok": stale_ok,
                "current_root_ok": current_root_ok,
            }
        except Exception as exc:  # noqa: BLE001
            passed = False
            detail = {"reason": f"{type(exc).__name__}: {exc}"}
        per_task.append(TaskResult(task_id=task_id, passed=passed, detail=detail))

    total = len(per_task)
    pres_rate = sum(1 for t in per_task if t.passed) / total if total else 0.0
    meta = fixture_meta(fixtures_path)
    return _build_result(
        meta.get("sub_battery", "time"),
        fixtures_path,
        per_task,
        {"temporal_context_preservation_rate": pres_rate},
    )


# Triangulate / Timing alias names from Phase 1a stubs are removed —
# canonical Dav1DPrometheus names below.

SUB_BATTERIES = {
    "transfer": run_transfer,                    # legacy SQD name (Phase 1a)
    "transfer-learning": run_transfer_learning,  # canonical Dav1DPrometheus
    "triangulation": run_triangulation,
    "truthtables": run_truthtables,
    "transitivity": run_transitivity,
    "time": run_time,
}