iolai26-solve/solver/router.py

"""Router: dispatch each puzzle to its symbolic solver (numerals, matching,
tables, translation) and return one answer per item. Never raises or empties."""

from __future__ import annotations

import re
from typing import Callable, List, Optional, Tuple

from .budget import Budget
from .fallback import ensure_nonempty, fallback_answer
from .llm import LLMClient, NullClient
from .matching import solve_matching
from .numerals import extract_attested, induce
from .preprocess import Pair, Puzzle, QueryItem, normalize
from .synth import synthesize
from .tables import TableSolver
from .template import TemplateTranslator
from .verifier import evaluate, leave_one_out

_NUM_RX = re.compile(r"\d+")
_QUOTED = re.compile(r"[\"«]([^\"»]+)[\"»]|‘([^’]{2,})’|'([^']{2,})'")


def _payload(text: str) -> str:
    """Payload of a whole-query item: quoted material, else text after a
    colon, else the final word of an instruction-like sentence."""
    t = normalize(text)
    m = _QUOTED.search(t)
    if m:
        return next(g for g in m.groups() if g).strip()
    if ":" in t:
        tail = t.split(":", 1)[1].strip()
        if tail:
            return tail
    m2 = re.match(r"^(give|translate|write|say|transcribe)\b.*\b(?:word|numeral|phrase|form)\s+(\S+)\s*$",
                  t, re.IGNORECASE)
    if m2:
        return m2.group(2).strip(".?!")
    return t


def solve_puzzle(puzzle: Puzzle, client: Optional[LLMClient] = None,
                 budget: Optional[Budget] = None, puzzles_left: int = 1) -> List[str]:
    return solve_puzzle_ex(puzzle, client, budget, puzzles_left)[0]


def solve_puzzle_ex(puzzle: Puzzle, client: Optional[LLMClient] = None,
                    budget: Optional[Budget] = None, puzzles_left: int = 1
                    ) -> Tuple[List[str], List[float], List[str]]:
    """Returns (answers, confidences, methods). Confidence is the verifier
    evidence behind each answer (LOO/eval fit of the solver that produced it,
    or ~1.0 for round-trip-verified numeral systems); 0.0 marks answers that
    came from the never-empty fallback ladder — those are the items worth LLM
    budget. Methods name the producing solver, for prompt candidate blocks
    and explanation-track traces."""
    client = client or NullClient()
    budget = budget or Budget()
    items = puzzle.items or [QueryItem(number="", text=puzzle.query or "")]
    try:
        answers, confs, methods = _dispatch(puzzle, items, client, budget, puzzles_left)
    except Exception:
        answers, confs, methods = None, None, None
    if answers is None:
        answers = [None] * len(items)
    if confs is None:
        confs = [0.0] * len(items)
    if methods is None:
        methods = ["none"] * len(items)
    answers = (list(answers) + [None] * len(items))[: len(items)]
    confs = (list(confs) + [0.0] * len(items))[: len(items)]
    methods = (list(methods) + ["none"] * len(items))[: len(items)]
    out = []
    for i, (item, ans) in enumerate(zip(items, answers)):
        if ans is None or not str(ans).strip():
            confs[i] = 0.0
            methods[i] = "fallback"
        direction = item.direction or "to_work"
        out.append(ensure_nonempty(ans, _payload(item.text), puzzle.pairs, direction))
    return out, confs, methods


def _dispatch(puzzle: Puzzle, items: List[QueryItem], client: LLMClient,
              budget: Budget, left: int
              ) -> Tuple[List[Optional[str]], List[float], List[str]]:
    tt = puzzle.task_type
    if tt in ("text_to_num", "num_to_text"):
        return _solve_numerals(puzzle, items)
    if tt == "match_letters":
        return _solve_matching(puzzle, items)

    table = TableSolver(puzzle)
    answers: List[Optional[str]] = [None] * len(items)
    confs: List[float] = [0.0] * len(items)
    methods: List[str] = ["none"] * len(items)
    plain_idx = []
    for i, it in enumerate(items):
        if it.row is not None and table.usable:
            ans_conf = table.solve(it)
            if ans_conf is not None:
                answers[i], confs[i] = ans_conf
                methods[i] = "table completion"
        if answers[i] is None:
            plain_idx.append(i)
    if plain_idx:
        translated, t_confs, t_methods = _solve_translation(
            puzzle, [items[i] for i in plain_idx], client, budget, left)
        for i, ans, c, m in zip(plain_idx, translated, t_confs, t_methods):
            answers[i], confs[i], methods[i] = ans, c, m
    return answers, confs, methods


# ---------------------------------------------------------------- numerals

def _solve_numerals(puzzle: Puzzle, items: List[QueryItem]
                    ) -> Tuple[List[Optional[str]], List[float]]:
    attested = extract_attested(puzzle.pairs)
    system = induce(attested) if attested else None
    out: List[Optional[str]] = []
    for item in items:
        text = _payload(item.text)
        if puzzle.task_type == "text_to_num":
            val = system.text_to_num(text) if system else None
            out.append(str(val) if val is not None else None)
        else:
            m = _NUM_RX.search(text)
            if system and m:
                out.append(system.num_to_text(int(m.group(0)), [p for p, _ in attested]))
            else:
                out.append(None)
    # an induced system is round-trip verified on every attested equation
    confs = [0.9 if a is not None else 0.0 for a in out]
    methods = ["numeral system (verified)" if a is not None else "none" for a in out]
    return out, confs, methods


# ---------------------------------------------------------------- matching

def _solve_matching(puzzle: Puzzle, items: List[QueryItem]
                    ) -> Tuple[List[Optional[str]], List[float]]:
    """Items are forms; options are the lettered meanings (from context or
    query). Answers are option letters when options exist, else the matched
    meaning text."""
    forms = [_payload(it.text) for it in items]
    if puzzle.lettered:
        labels = sorted(puzzle.lettered)
        meanings = [puzzle.lettered[l] for l in labels]
    else:
        labels = None
        meanings = [p.tgt for p in puzzle.pairs]
    if not meanings:
        return [None] * len(forms), [0.0] * len(forms), ["none"] * len(forms)
    matched = dict(solve_matching(forms, meanings, puzzle.pairs))
    out: List[Optional[str]] = []
    for f in forms:
        m = matched.get(f)
        if m and labels:
            out.append(labels[meanings.index(m)])
        else:
            out.append(m)
    # Hungarian is optimal for its score matrix, but the matrix itself is only
    # as good as the alignment evidence behind it — real-data EM is low, so
    # this stays BELOW the pipeline's keep-threshold: it surfaces as a
    # candidate hint in the LLM prompt rather than overriding the LLM
    conf = 0.45 if puzzle.pairs else 0.25
    return (out, [conf if a else 0.0 for a in out],
            ["optimal matching" if a else "none" for a in out])


# ------------------------------------------------------------- translation

def _solve_translation(puzzle: Puzzle, items: List[QueryItem], client: LLMClient,
                       budget: Budget, left: int
                       ) -> Tuple[List[Optional[str]], List[float]]:
    directions = {it.direction or "to_work" for it in items}
    primary = "to_task" if "to_task" in directions else "to_work"
    rounds = budget.cegis_rounds(left) if budget.allow_llm(left) else -1
    synth_res = synthesize(puzzle, client, primary, rounds) if rounds >= 0 else None

    solvers = {d: _pick_direction_solver(puzzle, synth_res, d) for d in directions}
    answers, confs, methods = [], [], []
    for it in items:
        ans, conf, method = solvers[it.direction or "to_work"](_payload(it.text))
        answers.append(ans)
        confs.append(conf)
        methods.append(method)
    return answers, confs, methods


def _pick_direction_solver(puzzle: Puzzle, synth_res, d: str) -> Callable[[str], Optional[str]]:
    """Rank candidate solvers honestly and chain them (first non-None answer
    wins). The grammar is a fixed program so it is evaluated directly on the
    attested pairs (it must reproduce them); the template translator and the
    fallback are *fit from* those pairs (they memorize them), so they are
    scored leave-one-out — otherwise memorization would always beat a
    generalizing grammar."""
    attested = [(p.tgt, p.src) if d == "to_task" else (p.src, p.tgt) for p in puzzle.pairs]

    def _subset(held_in_pairs) -> List[Pair]:
        keep = set(held_in_pairs)
        return [p for p in puzzle.pairs
                if ((p.tgt, p.src) if d == "to_task" else (p.src, p.tgt)) in keep]

    ranked: List[Tuple[float, int, str, Callable[[str], Optional[str]]]] = []

    if synth_res and synth_res.interpreter:
        fn = synth_res.interpreter.generate if d == "to_task" else synth_res.interpreter.analyze
        v = evaluate(fn, attested, synth_res.grammar.mdl())
        ranked.append((v.score, 0, "induced grammar", fn))

    tmpl = TemplateTranslator(puzzle.pairs, d)
    v_tmpl = leave_one_out(lambda held: TemplateTranslator(_subset(held), d).translate, attested)
    ranked.append((v_tmpl.score, 1, "template substitution", tmpl.translate))

    v_fb = leave_one_out(lambda held: (lambda q, kept=_subset(held): fallback_answer(q, kept, d)), attested)
    ranked.append((v_fb.score, 2, "nearest attested", lambda q: fallback_answer(q, puzzle.pairs, d)))

    ranked.sort(key=lambda t: (-t[0], t[1]))

    def solve(q: str) -> Tuple[Optional[str], float, str]:
        for score, _, name, fn in ranked:
            ans = fn(q)
            if ans:
                return ans, max(score, 0.0), name
        return None, 0.0, "none"

    return solve