app/fp_retrieve.py

"""Family Practice (00) hybrid retrieval + Schedule relationship graph.

Lexical (+ optional semantic) search is constrained to the Family Practice &
Practice in General (00) specialty set so retrieval does not drift into other
specialty sections (e.g. Paediatrics 26).
"""
from __future__ import annotations

import logging
from collections import defaultdict

from .config import settings
from .family_scope import (
    FAMILY_PRACTICE_SECTION,
    _normalize_doc,
    family_practice_00_filter,
    fetch_all_family_practice_00_primaries,
    filter_docs_family_practice_00_premium,
    is_family_practice_00_primary,
    is_family_practice_00_premium,
)
from .opensearch_client import get_client
from .paeds_retrieve import rrf_fuse
from .search import PRIMARY_ONLY, get_code
from .specialty_scope import filter_catalog_by_prefixes
from .weave_trace import redact_inputs, summarize_docs, traced

logger = logging.getLogger(__name__)

_RRF_K = 60

# Setting mirrors for common FP assessments / consults.
SETTING_MIRROR_GROUPS: dict[str, tuple[str, ...]] = {
    "general_assessment": ("A003", "C003"),
    "general_reassessment": ("A004", "C004", "W004"),
    # Office A006 ↔ hospital C005/C006 ↔ LTC W105/W106 when setting matches.
    "consultation": ("A006", "C005", "C006", "W105", "W106"),
    "hospital_subsequent": ("C002", "C007", "C009"),
    "ltc_subsequent": ("W001", "W002", "W003", "W008"),
}

# Content ladder (office) — not a strict consult ladder like Paediatrics.
ASSESSMENT_LADDER: tuple[str, ...] = (
    "mini_assessment",
    "minor_assessment",
    "intermediate_assessment",
    "general_assessment",
    "general_reassessment",
)

_FAMILY_PREMIUMS: dict[str, tuple[str, ...]] = {
    "office_assessment": (
        "E078",
        "E080",
        "G372",
        "G373",
        "G590",
        "E402",
        "E403",
        "E409",
        "E410",
        "A110",
        "A112",
    ),
    "well_baby": ("G372", "G373", "G590", "E402", "E403"),
    "consultation": ("E078", "E402", "E403", "E409", "E410"),
    "hospital_inpatient": (
        "C102",
        "C103",
        "C104",
        "C105",
        "C106",
        "C107",
        "C108",
        "C109",
        "C110",
        "E078",
        "E409",
        "E410",
    ),
    "ltc": ("E078", "E402", "E403"),
    "periodic_health": ("E078", "G372", "G373"),
}

_CODE_TO_FAMILY: dict[str, str] = {}
for code in ("A001", "A007", "A008", "A003", "A004"):
    _CODE_TO_FAMILY[code] = "office_assessment"
_CODE_TO_FAMILY["A002"] = "well_baby"
for code in ("A006", "C005", "C006", "W105", "W106"):
    _CODE_TO_FAMILY[code] = "consultation"
for code in ("C002", "C003", "C004", "C007", "C008", "C009", "C010", "C882"):
    _CODE_TO_FAMILY[code] = "hospital_inpatient"
for code in ("W001", "W002", "W003", "W004", "W008", "W010", "W109"):
    _CODE_TO_FAMILY[code] = "ltc"
_CODE_TO_FAMILY["K017"] = "periodic_health"

_CODE_TO_MIRRORS: dict[str, tuple[str, ...]] = {}
for _group, members in SETTING_MIRROR_GROUPS.items():
    for m in members:
        _CODE_TO_MIRRORS[m] = tuple(x for x in members if x != m)

_CODE_TO_GROUP: dict[str, str] = {
    m: g for g, members in SETTING_MIRROR_GROUPS.items() for m in members
}


def relationship_meta(code: str) -> dict:
    code = (code or "").upper()
    family = _CODE_TO_FAMILY.get(code)
    mirrors = list(_CODE_TO_MIRRORS.get(code, ()))
    group = _CODE_TO_GROUP.get(code)
    related_premiums = list(_FAMILY_PREMIUMS.get(family, ())) if family else []
    ladder_note = None
    if family == "office_assessment":
        ladder_note = (
            "Office assessment ladder: mini A008 → minor A001 → intermediate "
            "A007 → general A003 → re-assessment A004 / periodic health K017"
        )
    return {
        "service_family": family,
        "setting_mirrors": mirrors,
        "mirror_group": group,
        "related_premiums": related_premiums,
        "consult_ladder": ladder_note,
        "section_alternatives": [],
    }


def enrich_with_relationships(doc: dict) -> dict:
    out = dict(doc)
    meta = relationship_meta(out.get("billing_code") or "")
    out["relationship_family"] = meta["service_family"]
    out["relationship_mirrors"] = meta["setting_mirrors"]
    out["relationship_premiums"] = meta["related_premiums"]
    out["relationship_ladder"] = meta["consult_ladder"]
    out["relationship_section_alternatives"] = meta["section_alternatives"]
    return out


def related_premium_codes_for_primaries(primary_codes: list[str]) -> list[str]:
    """Premiums tied to each primary's service family only (no global dump).

    Hospital special-visit matrix codes (C102–C110) stay with hospital_inpatient
    primaries; office visits do not inherit them via allowlist bleed.
    """
    out: list[str] = []
    seen: set[str] = set()
    for code in primary_codes:
        family = _CODE_TO_FAMILY.get(code.upper())
        for prem in _FAMILY_PREMIUMS.get(family, ()):
            if prem not in seen:
                seen.add(prem)
                out.append(prem)
    return out


def _hits_to_docs(hits: list[dict]) -> list[dict]:
    docs = []
    for h in hits:
        src = h["_source"]
        doc = _normalize_doc(src, score=h.get("_score"))
        if is_family_practice_00_primary(doc):
            docs.append(enrich_with_relationships(doc))
    return docs


def _bm25_fp(query_text: str, filters: list[dict], top_k: int) -> list[dict]:
    client = get_client()
    body = {
        "size": top_k,
        "_source": {"excludes": ["code_vector"]},
        "query": {
            "bool": {
                "must": [
                    {
                        "multi_match": {
                            "query": query_text,
                            "fields": [
                                "billing_code^3",
                                "description_text^2",
                                "rules_and_constraints",
                                "parent_section",
                            ],
                        }
                    }
                ],
                "filter": filters,
            }
        },
    }
    response = client.search(index=settings.opensearch_index, body=body)
    return _hits_to_docs(response.get("hits", {}).get("hits", []))


def _knn_fp(
    query_vector: list[float] | None, filters: list[dict], top_k: int
) -> list[dict]:
    if not query_vector:
        return []
    client = get_client()
    body = {
        "size": top_k,
        "_source": {"excludes": ["code_vector"]},
        "query": {
            "knn": {
                "code_vector": {
                    "vector": query_vector,
                    "k": top_k,
                    "filter": {"bool": {"filter": filters}},
                }
            }
        },
    }
    try:
        response = client.search(index=settings.opensearch_index, body=body)
    except Exception:  # noqa: BLE001
        logger.exception("Family Practice k-NN search failed; continuing BM25 only")
        return []
    return _hits_to_docs(response.get("hits", {}).get("hits", []))


@traced(
    "hybrid_search_family_practice",
    kind="search",
    postprocess_inputs=redact_inputs,
    postprocess_output=summarize_docs,
)
def hybrid_search_family_practice(
    query_text: str,
    query_vector: list[float] | None,
    *,
    prefixes: list[str] | None = None,
    top_k: int | None = None,
) -> list[dict]:
    k = top_k or max(settings.hybrid_top_k, 40)
    filters = PRIMARY_ONLY + family_practice_00_filter()
    if prefixes:
        from .encounter_align import prefix_filter  # noqa: PLC0415

        pref = prefix_filter(prefixes)
        lexical = _bm25_fp(query_text, filters + pref, k)
        semantic = _knn_fp(query_vector, filters + pref, k)
        if not lexical and not semantic:
            lexical = _bm25_fp(query_text, filters, k)
            semantic = _knn_fp(query_vector, filters, k)
    else:
        lexical = _bm25_fp(query_text, filters, k)
        semantic = _knn_fp(query_vector, filters, k)

    fused = rrf_fuse([lexical, semantic]) if semantic else lexical
    logger.info(
        "Family Practice hybrid: lexical=%d semantic=%d fused=%d section=%s",
        len(lexical),
        len(semantic),
        len(fused),
        FAMILY_PRACTICE_SECTION,
    )
    return fused


def rank_full_fp_catalog(
    query_text: str,
    query_vector: list[float] | None,
    *,
    prefixes: list[str] | None = None,
) -> list[dict]:
    """Hybrid-rank the Family Practice (00) catalog (prefix-scoped)."""
    catalog = fetch_all_family_practice_00_primaries()
    catalog = filter_catalog_by_prefixes(catalog, prefixes or [])
    if not catalog and prefixes:
        catalog = fetch_all_family_practice_00_primaries()

    top_k = max(len(catalog), settings.hybrid_top_k, 50)
    hybrid_hits = hybrid_search_family_practice(
        query_text,
        query_vector,
        prefixes=prefixes,
        top_k=top_k,
    )
    by_code = {d["billing_code"]: d for d in hybrid_hits}

    relationship_boost: list[dict] = []
    for hit in hybrid_hits[:12]:
        for mirror in hit.get("relationship_mirrors") or []:
            if prefixes and mirror[:1] not in {p.upper() for p in prefixes}:
                continue
            if mirror in by_code:
                continue
            doc = get_code(mirror)
            if doc and is_family_practice_00_primary(doc):
                relationship_boost.append(enrich_with_relationships(doc))

    fused = rrf_fuse(
        [hybrid_hits, relationship_boost] if relationship_boost else [hybrid_hits]
    )
    fused_codes = {d["billing_code"] for d in fused}
    for doc in catalog:
        code = doc["billing_code"]
        if code not in fused_codes:
            fused.append(enrich_with_relationships(doc))
            fused_codes.add(code)
    return fused


def hybrid_premium_search_family_practice(
    query_text: str,
    query_vector: list[float] | None,
    *,
    primary_codes: list[str],
    top_k: int = 24,
) -> list[dict]:
    from .search import premium_search  # noqa: PLC0415
    from .sob_grounding import filter_current_sob_only, is_current_sob  # noqa: PLC0415

    families = sorted(
        {_CODE_TO_FAMILY[c] for c in primary_codes if c in _CODE_TO_FAMILY}
    )
    family_kw = " ".join(f.replace("_", " ") for f in families)
    premium_query = (
        f"Family Practice specialty 00 {family_kw}. {query_text} "
        "premium add-on immunization injection special visit chronic disease "
        "OTHER PREMIUMS after hours procedure E409 E410"
    )
    hits = premium_search(premium_query, query_vector=query_vector, top_k=top_k)
    hits = filter_current_sob_only(hits)

    related = related_premium_codes_for_primaries(primary_codes)
    by_code = {h["billing_code"]: enrich_with_relationships(h) for h in hits}
    for code in related:
        if code in by_code:
            continue
        doc = get_code(code)
        if not is_current_sob(doc):
            continue
        if not is_family_practice_00_premium(doc):
            continue
        by_code[code] = enrich_with_relationships(doc)

    ordered: list[dict] = []
    seen: set[str] = set()
    for code in related:
        if code in by_code and code not in seen:
            ordered.append(by_code[code])
            seen.add(code)
    for code, doc in by_code.items():
        if code not in seen:
            ordered.append(doc)
            seen.add(code)
    return filter_docs_family_practice_00_premium(ordered)


# Aliases matching the Paediatrics retrieve surface used by analyze dispatch.
rank_full_specialty_catalog = rank_full_fp_catalog
hybrid_premium_search_specialty = hybrid_premium_search_family_practice