solver/template.py

"""Template translation by minimal-pair substitution — the strongest
zero-LLM translation baseline for constructed puzzles.

Idea: puzzles are built so query sentences differ from attested ones by a
small substitution. Find the attested pair whose source is closest to the
query (token-level), then replace the differing tokens in its *target* using
alignment links (align.py). Works in both directions. Also supports
morph-level substitution for single-word queries (paradigm cells).
"""

from __future__ import annotations

from collections import Counter
from typing import Dict, List, Optional, Tuple

from .align import align as build_align, one_to_one
from .preprocess import Pair, strip_punct, tokenize


def _toks(s: str) -> List[str]:
    return [strip_punct(t).casefold() for t in tokenize(s) if strip_punct(t)]


def _flip(pairs: List[Pair]) -> List[Pair]:
    return [Pair(src=p.tgt, tgt=p.src) for p in pairs]


class TemplateTranslator:
    """direction 'to_work': translate task->work; 'to_task': work->task."""

    def __init__(self, pairs: List[Pair], direction: str = "to_work"):
        self.pairs = pairs if direction == "to_work" else _flip(pairs)
        self.amap = build_align(self.pairs)  # src tok -> ranked [(tgt, score)]

    def _sub(self, src_tok: str) -> Optional[str]:
        cands = self.amap.get(src_tok.casefold())
        return cands[0][0] if cands else None

    def _sub_in(self, src_tok: str, pool: List[str]) -> Optional[str]:
        """Best candidate for src_tok that is present in pool (context-aware:
        a token may have both a bare and an inflected realization; the one
        actually in the template target is the right one)."""
        for c, _ in self.amap.get(src_tok.casefold(), []):
            if c in pool:
                return c
        return None

    def _sub_like(self, src_tok: str, model: str) -> Optional[str]:
        """Best candidate for src_tok, preferring one that shares an affix
        (prefix/suffix >= 2 chars) with `model` — the form it will replace.
        kupu:nakupu :: moko:namoko."""
        cands = self.amap.get(src_tok.casefold(), [])
        for c, _ in cands:
            if len(c) >= 2 and len(model) >= 2 and (c[:2] == model[:2] or c[-2:] == model[-2:]):
                return c
        return cands[0][0] if cands else None

    def translate(self, query: str) -> Optional[str]:
        q = _toks(query)
        if not q:
            return None
        # rank templates by token-bag distance, then by length mismatch: a
        # same-length template is a substitution frame; a much shorter one
        # (e.g. a single-word gloss) would force fabricating structure
        ranked = sorted(
            ((_bag_distance(q, _toks(p.src)), abs(len(_toks(p.src)) - len(q)),
              _toks(p.src), _toks(p.tgt)) for p in self.pairs),
            key=lambda x: (x[0], x[1]),
        )
        max_dist = max(2, len(q) // 2)
        for dist, _, s, t in ranked:
            if dist == 0:
                return " ".join(t)
            if dist > max_dist:
                break
            out = self._substitute(q, s, t)
            if out:
                return out
        return None

    def _substitute(self, q: List[str], s: List[str], t: List[str]) -> Optional[str]:
        """Swap the tokens where query and template source differ, mapping
        both sides through the alignment. Abstains (None) when any needed
        link is missing — a wrong-but-confident answer is worse than letting
        the next template or the fallback ladder take over."""
        q_extra = list((Counter(q) - Counter(s)).elements())
        s_extra = list((Counter(s) - Counter(q)).elements())
        out = list(t)
        used: set = set()
        for s_tok in s_extra:
            s_tgt = self._sub_in(s_tok, out)
            if s_tgt is None:
                return None
            repl = None
            for qi, q_tok in enumerate(q_extra):
                if qi in used:
                    continue
                q_tgt = self._sub_like(q_tok, s_tgt)
                if q_tgt:
                    repl = q_tgt
                    used.add(qi)
                    break
            if repl is None:
                return None
            out[out.index(s_tgt)] = repl
        for qi, q_tok in enumerate(q_extra):
            if qi not in used:
                q_tgt = self._sub(q_tok)
                if q_tgt and q_tgt not in out:
                    out.append(q_tgt)
        return " ".join(out) if out else None


def _bag_distance(a: List[str], b: List[str]) -> int:
    ca, cb = Counter(a), Counter(b)
    return sum((ca - cb).values()) + sum((cb - ca).values())


def paradigm_complete(stem: str, pairs: List[Pair], cue: str = "") -> Optional[str]:
    """Complete a paradigm cell: find attested form-pairs (a, b) sharing a
    stem, group them by their string edit, and apply the dominant edit to
    `stem`. `cue` (e.g. 'plural') restricts to pairs whose gloss relation
    mentions the cue when glosses are available."""
    from .analogy import edit_rules, apply_rule

    vocab: Dict[str, str] = {}  # form -> gloss
    for p in pairs:
        if " " not in p.src.strip():
            vocab[p.src.strip().casefold()] = p.tgt.strip().casefold()

    rules: Counter = Counter()
    for a in vocab:
        for b in vocab:
            if a != b and len(b) > len(a) and b.startswith(a[: max(2, len(a) - 1)]):
                for r in edit_rules(a, b):
                    if cue:
                        ga, gb = vocab.get(a, ""), vocab.get(b, "")
                        # cue must relate the two glosses (e.g. 'houses' vs 'house')
                        if not (ga and gb and (ga in gb or gb in ga)):
                            continue
                    rules[r] += 1
    for r, _ in rules.most_common(3):
        out = apply_rule(r, stem.casefold())
        if out and out != stem:
            return out
    return None