PTSRaiAgent/app/core/intake_questions.py

"""Question selection and exact versioned content, with no side effects."""

from dataclasses import dataclass

import yaml

from app.config import get_settings
from app.core.cssrs_content import load_cssrs_russian
from app.core.questionnaire import IntakePolicy, alcohol_expansion, score_alcohol, score_cssrs_recent
from app.core.scales import ScaleDef, load_scale


@dataclass(frozen=True)
class Question:
    id: str
    text: str
    options: tuple[tuple[object, str], ...] = ()
    period: str = "context"
    numeric: bool = False
    optional: bool = False
    maximum: float = 24


def content() -> dict:
    return yaml.safe_load((get_settings().config_dir / "instruments/intake_forms.yaml").read_text())


def texts() -> dict:
    return yaml.safe_load((get_settings().config_dir / "texts/intake.yaml").read_text())


def scale_question(scale_id: str, item_id: str) -> Question:
    forms = content()["forms"]
    scale = ScaleDef.model_validate(forms[scale_id]) if scale_id in forms else load_scale(scale_id)
    item = next(i for i in scale.items if i.id == item_id)
    instruction = scale.instruction or ""
    # The beverage table is part of the approved questionnaire source.
    if item_id == "rus_audit_2":
        instruction += "\n" + texts()["standard_drinks"]
    return Question(
        item_id,
        instruction + "\n\n" + item.text,
        tuple((o.value, o.label) for o in item.options),
        scale.recall_period or "",
    )


def group_answers(answers: dict, prefix: str) -> dict[str, int]:
    return {k: v for k, v in answers.items() if k.startswith(prefix + "_")}


def cssrs_answers(answers: dict) -> dict[str, int]:
    return {k.removeprefix("cssrs_"): v for k, v in answers.items() if k.startswith("cssrs_")}


def next_intake(answers: dict, policy: IntakePolicy, sex: str | None) -> Question | None:
    """Stage 1 in the fixed order of «Опросники для MVP»: PHQ-2, trauma filter → TSQ, RUS-AUDIT-S, C-SSRS screener,
    then the context questions; stage 2 (rest of PHQ-9, PCL-5, full RUS-AUDIT) only where stage 1 is positive.
    Items inside a scale keep the order of the source. The risk stop rules run after every answer (flows/intake.py),
    so a positive answer interrupts the questionnaire wherever it comes."""
    for item in ("phq9_1", "phq9_2"):
        if item not in answers:
            return scale_question("phq9", item)
    if "trauma" not in answers:
        return Question(
            "trauma", content()["trauma_introduction"] + "\n\n" + content()["trauma_filter"], ((1, "Да"), (0, "Нет"))
        )
    if answers["trauma"]:
        for item in load_scale("tsq").item_ids:
            if item not in answers:
                return scale_question("tsq", item)
        if "recent_trauma" not in answers:
            return Question("recent_trauma", texts()["recent_trauma"], ((1, "Да"), (0, "Нет")))
    for item in load_scale("rus_audit_s").item_ids:
        if item not in answers:
            return scale_question("rus_audit_s", item)
    c = score_cssrs_recent(cssrs_answers(answers))
    if c.missing:
        form = load_cssrs_russian()
        item = next(i for i in form.items if i.id == c.missing[0])
        period = texts()["periods"][item.recall_period]
        # Guidance is interviewer-facing; retained in the source and record version.
        return Question(
            "cssrs_" + item.id,
            period + "\n\n" + item.text,
            tuple((o.value, o.label) for o in form.options),
            item.recall_period,
        )
    if sex is None and "sex" not in answers:
        return Question("sex", texts()["sex"], (("f", "Женский"), ("m", "Мужской"), ("unknown", "Не указывать")))
    for key in ("tbi", "psychotropics"):
        if key not in answers:
            return Question(key, texts()[key], ((1, "Да"), (0, "Нет")))
    if answers["phq9_1"] + answers["phq9_2"] >= policy.phq2_expand_at:
        for item in load_scale("phq9").item_ids:
            if item not in answers:
                return scale_question("phq9", item)
    if answers["trauma"] and sum(group_answers(answers, "tsq").values()) >= policy.tsq_expand_at:
        for item in load_scale("pcl5").item_ids:
            if item not in answers:
                return scale_question("pcl5", item)
    alcohol = group_answers(answers, "rus_audit")
    selected_sex = sex or (answers.get("sex") if answers.get("sex") in ("m", "f") else None)
    short = score_alcohol(
        {k: v for k, v in alcohol.items() if k in load_scale("rus_audit_s").item_ids}, short=True, sex=selected_sex
    )
    if alcohol_expansion(short, policy):
        for item in load_scale("rus_audit").item_ids:
            if item == "rus_audit_2" and alcohol.get("rus_audit_1") == 0:
                continue
            if item not in answers:
                return scale_question("rus_audit", item)
    return None


def cssrs_since_contact_question(answers: dict) -> Question | None:
    """Daily suicide-risk module of stage 2: the C-SSRS screener asked «since the last contact».

    The official Russian Since Last Contact form is not in the repository yet (docs/clinical/README.md, C06), so the
    items of the Recent form are asked verbatim under the since-last-contact period. Behaviour since the last contact
    is recent by definition: the answer to q6 also fills q6_recent (derived), so the recent triage applies.
    """
    c = score_cssrs_recent(cssrs_answers(answers))
    missing = [k for k in c.missing if k != "q6_recent"]
    if not missing:
        return None
    form = load_cssrs_russian()
    item = next(i for i in form.items if i.id == missing[0])
    return Question(
        "cssrs_" + item.id,
        texts()["periods"]["since_last_contact"] + "\n\n" + item.text,
        tuple((o.value, o.label) for o in form.options),
        "since_last_contact",
    )


DIARY_PARTS = ("morning", "evening")
DIARY_KINDS = ("diary_morning", "diary_evening")


def diary_config() -> dict:
    return yaml.safe_load((get_settings().config_dir / "instruments/diary.yaml").read_text(encoding="utf-8"))


def diary_questions(part: str, modules: list[str], *, day_index: int = 0) -> tuple[Question, ...]:
    """Stage 2 diary of one part of the day: the core block plus, in the evening, the modules of the profile.

    PTSD («veteran») and depression run side by side when both are present. The depression module asks three
    PHQ-9 items a day in rotation by `day_index` (days since the first diary). Ids are prefixed (`dphq_`, `dpcl_`)
    so diary answers never mix with the full scales.
    """
    if part not in DIARY_PARTS:
        raise ValueError("unknown diary part")
    config = diary_config()
    questions = [
        Question(
            item["id"],
            item["text"],
            tuple((o["value"], o["label"]) for o in item.get("options", [])),
            "past_day",
            item.get("numeric", False),
            item.get("optional", False),
            item.get("max", 24),
        )
        for item in config[part]
    ]
    if part == "evening":
        forms = content()["forms"]
        if "depression" in modules:
            module = config["depression"]
            items = forms[module["form"]]["items"]
            per_day = module["items_per_day"]
            start = (day_index % (len(items) // per_day)) * per_day
            options = tuple((o["value"], o["label"]) for o in module["options"])
            for item in items[start : start + per_day]:
                number = item["id"].rsplit("_", 1)[1]
                questions.append(
                    Question(f"dphq_{number}", module["question"] + "\n\n" + item["text"], options, "today")
                )
        if "veteran" in modules:
            module = config["ptsd"]
            form = forms[module["form"]]
            by_id = {i["id"]: i["text"] for i in form["items"]}
            options = tuple((o["value"], o["label"]) for o in form["options"])
            for item_id in module["items"]:
                number = item_id.rsplit("_", 1)[1]
                questions.append(
                    Question(f"dpcl_{number}", module["question"] + "\n\n" + by_id[item_id], options, "today")
                )
    return tuple(questions)