Jyotisha/scripts/rectification/probe_question_contract.py

"""Shared probe → choice-card contract for Python event probes.

Must stay aligned with frontend/src/lib/rectification-agentic/v9/probe-question-contract.ts.
"""

from __future__ import annotations

import re
from typing import Any, Sequence

QUESTION_CONTRACT_VERSION = "probe-question-v1"
ANSWER_CLASSES = ("yes", "weak_yes", "no", "unsure")
EXISTENCE_STYLE_OPTIONS: tuple[dict[str, str], ...] = (
    {"label": "明确发生且时间吻合", "answer_class": "yes"},
    {"label": "发生过但程度较弱", "answer_class": "weak_yes"},
    {"label": "明确没有发生", "answer_class": "no"},
    {"label": "这段记不清楚", "answer_class": "unsure"},
)
QUALITY_STYLE_OPTIONS: tuple[dict[str, str], ...] = (
    {"label": "发挥明显失常或压力很大", "answer_class": "yes"},
    {"label": "有压力但不算明显失常", "answer_class": "weak_yes"},
    {"label": "发挥正常、没有明显失常", "answer_class": "no"},
    {"label": "这段记不清楚", "answer_class": "unsure"},
)
VARGA_NONE_STYLE_OPTION = {"label": "都不是这些特质", "answer_class": "no"}
UNSURE_STYLE_OPTION = {"label": "这段记不清楚", "answer_class": "unsure"}
_FORBIDDEN = ("外貌", "体质", "胎记", "疤痕", "伤疤", "身高", "体型")
_CLOCK = re.compile(r"(?:[01]?\d|2[0-3]):[0-5]\d")


def probe_question_kind(value: Any) -> str:
    if value in {"varga_style", "event_quality"}:
        return str(value)
    return "existence"


def clipped_probe_label(value: Any, minimum: int = 4, maximum: int = 80) -> str | None:
    if not isinstance(value, str):
        return None
    text = " ".join(value.split())
    if len(text) < minimum or len(text) > maximum:
        return None
    if any(token in text for token in _FORBIDDEN) or _CLOCK.search(text):
        return None
    return text


def _incoming_option(row: Any) -> dict[str, str] | None:
    if not isinstance(row, dict):
        return None
    answer = row.get("answer_class") or row.get("answerClass")
    label = clipped_probe_label(row.get("label"))
    if not label or answer not in ANSWER_CLASSES:
        return None
    payload = {"label": label, "answer_class": str(answer)}
    sign = row.get("sign")
    if isinstance(sign, str) and sign.strip():
        payload["sign"] = sign.strip()
    return payload


def complete_style_options(
    choice_kind: Any,
    style_options: Sequence[Any] | None = None,
) -> list[dict[str, str]] | None:
    kind = probe_question_kind(choice_kind)
    incoming = [item for item in (_incoming_option(row) for row in (style_options or [])) if item]
    by_class: dict[str, dict[str, str]] = {}
    if kind == "varga_style":
        for option in incoming:
            by_class[option["answer_class"]] = option
        by_class.setdefault("unsure", dict(UNSURE_STYLE_OPTION))
        scoring = [item for item in ANSWER_CLASSES if item != "unsure" and item in by_class]
        if len(scoring) < 2:
            return None
        by_class.setdefault("no", dict(VARGA_NONE_STYLE_OPTION))
        if "yes" not in by_class or "weak_yes" not in by_class:
            return None
    else:
        catalog = QUALITY_STYLE_OPTIONS if kind == "event_quality" else EXISTENCE_STYLE_OPTIONS
        for option in catalog:
            by_class[option["answer_class"]] = dict(option)
        for option in incoming:
            by_class[option["answer_class"]] = option
    ordered: list[dict[str, str]] = []
    seen: set[str] = set()
    for answer_class in ANSWER_CLASSES:
        option = by_class.get(answer_class)
        if not option:
            return None
        label = option["label"]
        if label in seen and option.get("sign"):
            label = f"{label}（{option['sign']}）"
        if label in seen:
            label = f"{label}·{answer_class}"
        seen.add(label)
        ordered.append({**option, "label": label})
    labels = {item["label"] for item in ordered}
    classes = {item["answer_class"] for item in ordered}
    if len(ordered) != 4 or labels != {item["label"] for item in ordered} or classes != set(ANSWER_CLASSES):
        return None
    if len(labels) != 4:
        return None
    return ordered


def is_renderable_probe(probe: dict[str, Any]) -> bool:
    gain = probe.get("information_gain")
    if not isinstance(gain, (int, float)) or gain <= 0:
        return False
    candidate_ids = probe.get("candidate_ids") or []
    outcomes = probe.get("expected_outcomes") or []
    if len(candidate_ids) < 2 or len(outcomes) < 2:
        return False
    return complete_style_options(probe.get("choice_kind"), probe.get("style_options")) is not None