import re

from agents.variants.base import TeacherVariant
from agents.verify import answer_matches

ANSWER_ONLY = re.compile(r"^\s*The correct final answer is\s*\\boxed\{.+\}\.?\s*$", re.S)


class P2b(TeacherVariant):
    """Flywheel bet: identical to P2 except the truncated branch, which gets a tight
    in-voice full solution to pressure the student to finish within the token cap.
    Kill-switch: retire if avg@k dips while completion-rate rises."""
    name = "p2b_flywheel"
    description = "= P2, but truncated rollouts get a tight full solution (compression bet)"

    def accepts(self, prefix, reference):
        if ANSWER_ONLY.match(prefix or ""):
            return answer_matches(prefix, reference)
        return super().accepts(prefix, reference)

    def compliant(self, r, attempt, tok):
        """Truncated-branch compliance: r must be <= 0.8x the attempt in tokens."""
        lr = len(tok(r or "", add_special_tokens=False)["input_ids"])
        la = len(tok(attempt or "", add_special_tokens=False)["input_ids"]) or 1
        return lr <= 0.8 * la


VARIANT = P2b