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


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:
    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,
        )
    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)
    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 diary_questions(
    modules: list[str], *, suicide_monitoring: bool, weekday: int, day_index: int | None = None
) -> tuple[Question, ...]:
    config = yaml.safe_load((get_settings().config_dir / "instruments/diary.yaml").read_text())
    selected = ["core"]
    detailed = [name for name in ("veteran", "depression") if name in modules]
    if len(detailed) == 2:
        detailed = [detailed[(day_index if day_index is not None else weekday) % 2]]
    if "veteran" in detailed:
        selected.append("ptsd")
    if "depression" in detailed:
        selected.append("depression")
    if suicide_monitoring or "depression" in modules:
        selected.append("suicide")
    questions = []
    for section in selected:
        for item in config[section]:
            if item.get("weekly") and weekday != 0:
                continue
            questions.append(
                Question(
                    item["id"],
                    item["text"],
                    tuple((o["value"], o["label"]) for o in item.get("options", [])),
                    "past_week" if item.get("weekly") else "past_day",
                    item.get("numeric", False),
                    item.get("optional", False),
                )
            )
    questions.sort(key=lambda q: q.id != "death_thoughts")
    return tuple(questions)