src/gemeo/ground_sus.py

"""SUS grounding — PCDT/CEAF/UF-aware compliance check.

For each top-1 diagnosis, returns:
  - Whether a PCDT exists (`brazilian_context.get_pcdt`)
  - Recommended therapies
  - Whether the patient's UF dispenses each therapy (heuristic until APAC pulled)
  - Nearest centro de referência
  - Triagem neonatal coverage
  - Patient associations

This is the unique differentiator vs. DeepRare/global tools — every rec is
grounded in what SUS actually delivers in a specific UF.
"""
from __future__ import annotations
import logging
from typing import Optional

from .types import SusCheck

logger = logging.getLogger("gemeo.ground_sus")


def _safe_get_pcdt(orpha: str):
    try:
        from brazilian_context import get_pcdt
        return get_pcdt(orpha)
    except Exception:
        return None


def _safe_get_centros(uf: Optional[str], specialty: Optional[str] = None):
    try:
        from brazilian_context import get_centros_referencia
        return get_centros_referencia(uf, specialty) or []
    except Exception:
        return []


def _safe_get_triagem(orpha: Optional[str]):
    try:
        from brazilian_context import get_triagem_neonatal
        result = get_triagem_neonatal(orpha)
        if isinstance(result, dict):
            return result
        return {}
    except Exception:
        return {}


def _safe_get_associations(disease_name: Optional[str]):
    try:
        from brazilian_context import get_associacoes
        return get_associacoes(disease_name) or []
    except Exception:
        return []


def check(
    *,
    orpha: Optional[str],
    disease_name: Optional[str] = None,
    uf: Optional[str] = None,
) -> SusCheck:
    """Run the SUS grounding check for a (disease, UF) pair."""
    if not orpha:
        return SusCheck()

    pcdt = _safe_get_pcdt(orpha)
    has_pcdt = bool(pcdt)
    therapies = []
    pcdt_url = None
    if pcdt:
        pcdt_url = pcdt.get("url") or pcdt.get("link")
        for key in ("therapies", "medicamentos", "tratamento"):
            v = pcdt.get(key)
            if isinstance(v, list):
                therapies.extend(str(t) for t in v if t)
            elif isinstance(v, str):
                therapies.append(v)
    therapies = list(dict.fromkeys(therapies))[:10]

    # Dispensation heuristic — TODO: wire to APAC/SIA when ingested.
    # For now: assume CEAF therapies are universally dispensed if PCDT exists.
    # Override later with real per-UF data.
    dispensed = {t: bool(has_pcdt) for t in therapies}

    # Centros — query patient UF first, then fall back to nationwide.
    nearest = None
    centros = _safe_get_centros(uf)
    if not centros and uf:
        centros = _safe_get_centros(None)
    if centros:
        c = centros[0]
        nearest = {
            "nome": c.get("nome") or c.get("name"),
            "cidade": c.get("cidade") or c.get("city"),
            "uf": c.get("uf") or c.get("state"),
            "telefone": c.get("telefone") or c.get("phone"),
            "especialidade": c.get("especialidade") or c.get("specialty"),
        }

    triagem = _safe_get_triagem(orpha)
    triagem_includes = bool(triagem.get("included") or triagem.get("incluida") or triagem.get("teste_pezinho"))

    associations = _safe_get_associations(disease_name)
    associations_brief = [
        {
            "nome": a.get("nome") or a.get("name"),
            "site": a.get("site") or a.get("url"),
            "telefone": a.get("telefone"),
        }
        for a in associations[:5] if isinstance(a, dict)
    ]

    return SusCheck(
        disease_orpha=orpha,
        has_pcdt=has_pcdt,
        pcdt_url=pcdt_url,
        therapy_pcdt_recommended=therapies,
        therapy_dispensed_in_uf=dispensed,
        nearest_centro=nearest,
        triagem_neonatal_includes=triagem_includes,
        associations=associations_brief,
    )