app/specialty_scope.py

"""Scope recommendations to an OHIP provider specialty.

Today Paediatrics (26) is the only active specialty with full retrieval /
allow-list grounding. Multi-specialty expansion registers new profiles in
``specialty_registry`` and adds matching retrieve modules; this module remains
the Paediatrics (26) gate used by the analyze pipeline when that specialty is
selected.

Per the Schedule of Benefits / EPC guidance, specialists bill visit codes from
their specialty's Consultations and Visits listings — not Family Practice &
Practice in General (00). Primary recommendations for specialty 26 must come
from the Paediatrics specialty set in the **current** Schedule (2026), including
setting-equivalent listings (office / hospital / LTC / GENERAL LISTINGS mirrors).
"""
from __future__ import annotations

import logging
import re

from .config import settings
from .opensearch_client import get_client
from .sob_grounding import is_current_sob

logger = logging.getLogger(__name__)

PAEDIATRICS_SECTION = "PAEDIATRICS (26)"
PROVIDER_SPECIALTY_CODE = "26"
PROVIDER_SPECIALTY_LABEL = "26 - Paediatrics"

# Specialty visit / consult listings that live under other SoB section headers
# (GENERAL LISTINGS, hospital, LTC) but are Paediatrics (26) services — same
# conditions as the office Paediatrics listings (e.g. A260 ↔ C260).
_PAEDS_SETTING_EQUIVALENTS: frozenset[str] = frozenset(
    {
        # Office / GENERAL LISTINGS mirrors
        "A260",  # Special paediatric consultation
        "A265",  # Consultation
        "A662",  # Extended special paediatric consultation
        "A667",  # Neurodevelopmental consultation
        "A695",  # Neurodevelopmental consultation (parallel listing)
        "K122",  # Individual developmental/behavioural care
        # Hospital in-patient mirrors
        "C260",
        "C262",
        "C263",
        "C264",
        "C265",
        "C266",
        "C268",
        "C565",  # Limited consultation (mirror of A565)
        "C661",
        "C662",
        "C665",
        "C667",
        "C695",
        # Long-term care mirrors
        "W260",
        "W261",
        "W262",
        "W265",
        "W266",
        "W269",
        "W565",  # Limited consultation (mirror of A565)
        "W662",
        "W667",
        "W695",
    }
)

# False positives from naive A26/H26 prefix matching (not Paediatrics specialty).
_PAEDS_PREFIX_EXCLUDE: frozenset[str] = frozenset(
    {
        "H264",  # ER pelvic exam — Family Practice listing
    }
)

_PAEDS_DESC = re.compile(
    r"paediatric|pediatric|neurodevelopmental|well baby|newborn care",
    re.I,
)

# Premium / add-on sections a Paediatrics (26) physician may claim in addition
# to specialty visit codes (current Schedule). Used to keep after-hours etc.
_PAEDS_PREMIUM_SECTION_HINTS: tuple[str, ...] = (
    "AFTER HOURS SPECIAL VISIT",
    "SPECIAL VISIT PREMIUM",
    "OTHER PREMIUMS",  # after-hours procedure % premiums (E409/E410)
    "CHRONIC DISEASE ASSESSMENT PREMIUM",
    "FIRST VISIT BY PRIMARY CARE PHYSICIAN AFTER HOSPITAL DISCHARGE",
    "INTRAMUSCULAR, SUBCUTANEOUS OR INTRADERMAL",
    "IMMUNIZATION",
)

# Visit-care after-hours / special-visit / common paeds add-ons (current SoB).
# Do NOT include CT/MRI interpretation premiums (E406/E408 etc.) or ED-only
# procedure premiums (E412/E413).
_PAEDS_PREMIUM_CODE_ALLOW: frozenset[str] = frozenset(
    {
        # After Hours Special Visit Premiums (visit care — Paediatrics-compatible)
        "E402",  # evenings Mon–Fri / weekend daytime
        "E403",  # nights
        # OTHER PREMIUMS — after-hours procedure % (non-ED physician)
        "E409",  # evenings / weekend-holiday daytime+evenings — +50% procedural
        "E410",  # nights — +75% procedural
        # Chronic disease (Schedule lists specialty 26)
        "E078",
        # Common paediatric add-ons
        "G372",
        "G373",
        "G590",
        # Special-visit matrix (first/additional/travel × time band)
        "C102",
        "C103",
        "C104",
        "C105",
        "C106",
        "C107",
        "C108",
        "C109",
        "C110",
    }
)

# Not Paediatrics-visit-compatible: ED-only procedure premiums and CT/MRI
# interpretation premiums. E409/E410 (OTHER PREMIUMS, non-ED) are allowlisted.
_PAEDS_PREMIUM_CODE_EXCLUDE: frozenset[str] = frozenset(
    {"E412", "E413", "E406", "E407", "E408"}
)

_DOC_SOURCE_EXCLUDES = ["code_vector"]


def paeds_premium_allowlist_codes() -> frozenset[str]:
    """Known Paediatrics-compatible premium codes to inject when BM25 misses them."""
    return _PAEDS_PREMIUM_CODE_ALLOW


def paeds_setting_equivalent_codes() -> frozenset[str]:
    return _PAEDS_SETTING_EQUIVALENTS


def _is_imaging_after_hours(doc: dict) -> bool:
    section = (doc.get("parent_section") or "").upper()
    desc = (doc.get("description_text") or doc.get("description") or "").lower()
    if "CT/MRI" in section or "CT / MRI" in section:
        return True
    if "ct/mri" in desc or "ct / mri" in desc:
        return True
    if "interpretation" in desc and ("ct" in desc or "mri" in desc):
        return True
    return False


def _normalize_doc(src: dict, score: float | None = None) -> dict:
    return {
        "billing_code": src["billing_code"],
        "description_text": src.get("description_text", ""),
        "rules_and_constraints": src.get("rules_and_constraints", ""),
        "base_fee_cad": src.get("base_fee_cad", 0.0),
        "parent_section": src.get("parent_section", ""),
        "reference": src.get("reference"),
        "differentiators": src.get("differentiators"),
        "effective_date": src.get("effective_date"),
        "termination_date": src.get("termination_date"),
        "fee_components": src.get("fee_components"),
        "in_current_schedule": src.get("in_current_schedule") is True,
        "score": score,
    }


def is_paediatrics_26_primary(doc: dict) -> bool:
    """True if this code is a Paediatrics (26) specialty visit/consult listing."""
    code = (doc.get("billing_code") or "").upper()
    if not code or code in _PAEDS_PREFIX_EXCLUDE:
        return False
    if not is_current_sob(doc):
        return False
    section = (doc.get("parent_section") or "").strip()
    if section == PAEDIATRICS_SECTION:
        return True
    if code in _PAEDS_SETTING_EQUIVALENTS:
        return True
    # Catch additional setting mirrors whose description ties them to paeds
    # specialty consults (e.g. "subject to the same conditions as A260").
    desc = doc.get("description_text") or doc.get("description") or ""
    if _PAEDS_DESC.search(desc) and code[0] in "ACWHK":
        # Exclude Family Practice (00) general well-baby etc. — those are for
        # specialty 00; Paediatrics uses A268 / A261 / A262 instead.
        if "FAMILY PRACTICE" in section.upper():
            return False
        return True
    return False


def is_paediatrics_26_premium(doc: dict) -> bool:
    """True if this premium/add-on is claimable alongside Paediatrics (26) care.

    Strictly limited to codes listed in the **most recent** Schedule of Benefits
    (``in_current_schedule is True``). Legacy FSM-only fee lines never qualify.
    Includes OTHER PREMIUMS E409/E410 (non-ED after-hours procedure %). Excludes
    ED-only E412/E413 and CT/MRI interpretation premiums (E406–E408).
    """
    code = (doc.get("billing_code") or "").upper()
    if not code:
        return False
    if code in _PAEDS_PREMIUM_CODE_EXCLUDE:
        return False
    if not is_current_sob(doc):
        return False
    if _is_imaging_after_hours(doc):
        return False
    # Specialty-section age/developmental premiums (K267/K269/K119…)
    if is_paediatrics_26_primary(doc):
        return True
    if code in _PAEDS_PREMIUM_CODE_ALLOW:
        return True
    section = (doc.get("parent_section") or "").upper()
    if any(h in section for h in _PAEDS_PREMIUM_SECTION_HINTS):
        return True
    return False


def paediatrics_26_filter() -> list[dict]:
    """OpenSearch filter restricting to Paediatrics (26) specialty visit codes."""
    return [
        {
            "bool": {
                "should": [
                    {"term": {"parent_section": PAEDIATRICS_SECTION}},
                    {"terms": {"billing_code": sorted(_PAEDS_SETTING_EQUIVALENTS)}},
                ],
                "minimum_should_match": 1,
            }
        }
    ]


def fetch_all_paediatrics_26_primaries() -> list[dict]:
    """Load every current-schedule Paediatrics (26) primary code from the index.

    Includes the full ``PAEDIATRICS (26)`` section plus setting-equivalent
    mirrors so the LLM always sees the complete specialty catalog (not only
    BM25 top-K hits).
    """
    client = get_client()
    body = {
        "size": 200,
        "_source": {"excludes": _DOC_SOURCE_EXCLUDES},
        "query": {
            "bool": {
                "filter": [
                    {"term": {"in_current_schedule": True}},
                    {
                        "bool": {
                            "should": [
                                {"term": {"parent_section": PAEDIATRICS_SECTION}},
                                {
                                    "terms": {
                                        "billing_code": sorted(
                                            _PAEDS_SETTING_EQUIVALENTS
                                        )
                                    }
                                },
                            ],
                            "minimum_should_match": 1,
                        }
                    },
                ]
            }
        },
    }
    try:
        response = client.search(index=settings.opensearch_index, body=body)
    except Exception:  # noqa: BLE001
        logger.exception("Failed to fetch Paediatrics (26) catalog")
        return []

    docs: list[dict] = []
    seen: set[str] = set()
    for hit in response.get("hits", {}).get("hits", []):
        src = hit["_source"]
        code = src.get("billing_code")
        if not code or code in seen:
            continue
        doc = _normalize_doc(src)
        if not is_paediatrics_26_primary(doc):
            continue
        seen.add(code)
        docs.append(doc)
    docs.sort(key=lambda d: d["billing_code"])
    return docs


def filter_catalog_by_prefixes(
    catalog: list[dict], prefixes: list[str]
) -> list[dict]:
    """Keep catalog codes whose billing_code starts with one of ``prefixes``."""
    if not prefixes:
        return list(catalog)
    allowed = {p.upper() for p in prefixes}
    return [d for d in catalog if (d.get("billing_code") or "")[:1] in allowed]


def merge_paeds_catalog_into_context(
    retrieved: list[dict],
    catalog: list[dict],
) -> list[dict]:
    """Ensure the full (prefix-scoped) Paediatrics catalog is in LLM context.

    Retrieved hits keep their ranking first; any catalog codes BM25 missed are
    appended so every relevant Paediatrics (26) listing is considered.
    """
    merged: list[dict] = []
    seen: set[str] = set()
    for doc in retrieved + catalog:
        code = doc.get("billing_code")
        if not code or code in seen:
            continue
        if not is_paediatrics_26_primary(doc):
            continue
        if doc.get("in_current_schedule") is False:
            continue
        seen.add(code)
        merged.append(doc)
    return merged


def filter_docs_paediatrics_26_primary(docs: list[dict]) -> list[dict]:
    return [d for d in docs if is_paediatrics_26_primary(d)]


def filter_docs_paediatrics_26_premium(docs: list[dict]) -> list[dict]:
    return [d for d in docs if is_paediatrics_26_premium(d)]