app/deterministic_rank.py

"""Deterministic primary + premium selection for clinic-pilot consistency.

Retrieval order and Schedule rules pick the codes. The LLM may only write
justifications afterward — it must not choose or reorder the set.
"""
from __future__ import annotations

import re
from typing import Any

from collections.abc import Callable

from .age_align import AgeAlignment, premium_applies_to_age
from .config import settings
from .schemas import AssociatedCode, CaseAnalysisResponse, CodeRecommendation
from .sob_grounding import is_current_sob, sob_grounded_rules
from .specialty_scope import is_paediatrics_26_premium, is_paediatrics_26_primary
from .time_align import premium_applies_to_band, premium_time_bands

PrimaryPred = Callable[[dict], bool]
PremiumPred = Callable[[dict], bool]

# Rank → fixed confidence bands (no LLM sampling).
_PRIMARY_CI = ("90-95%", "80-85%", "70-75%", "65-70%", "60-65%")
_PREMIUM_CI = ("85-90%", "75-80%", "70-75%", "65-70%", "60-65%")

_CHRONIC_RE = re.compile(
    r"\bchronic\b|\basthma\b|\bdiabet|\bCOPD\b|\bepilep|\bseizure disorder\b|"
    r"\bcongenital\b|\bcystic fibrosis\b|\bcerebral palsy\b",
    re.I,
)
_IMMUNIZ_RE = re.compile(
    r"\bimmuni[sz]|\bvaccine\b|\bvaccinat|\binjection\b|\bshot\b|\bbooster\b",
    re.I,
)
_PROCEDURE_RE = re.compile(
    r"\bprocedur|\bsutur|\blacerat|\bexcis|\bbiops|\bcast\b|\bsplin",
    re.I,
)
_DISCHARGE_RE = re.compile(
    r"\bdischarge\b|\bpost[\s-]?hospital\b|\bafter hospital\b|\brecent admission\b",
    re.I,
)

_IMMUNIZ_CODES = frozenset({"G372", "G373", "G590"})
_CHRONIC_CODES = frozenset({"E078"})
_DISCHARGE_CODES = frozenset({"E080"})
_PROCEDURE_PREMIUMS = frozenset({"E409", "E410"})
_AGE_ADDON_CODES = frozenset({"A110", "A112"})
_TIME_PREMIUM_HINT = re.compile(
    r"after[\s-]?hours|special visit|evening|night|weekend|holiday",
    re.I,
)

# Enhanced 18-month well-baby (FP A002 / Paediatrics A268) need explicit evidence.
_ENHANCED_18M_CODES = frozenset({"A002", "A268"})
_ENHANCED_18M_RE = re.compile(
    r"\benhanced\s+18|\b18[\s-]?month\b|\b18\s*mo(?:nth)?s?\b",
    re.I,
)
_WELL_BABY_RE = re.compile(r"\bwell[\s-]?baby\b|\bwell[\s-]?child\b", re.I)
_PERIODIC_RE = re.compile(
    r"\bperiodic\s+health\b|\bannual\s+health\b|\bcomplete\s+physical\b",
    re.I,
)
_LTC_RE = re.compile(
    r"\blong[\s-]?term\b|\bnursing\s+home\b|\bLTC\b|\bcontinuing\s+care\b",
    re.I,
)
_MINOR_HINT_RE = re.compile(
    r"\bminor\b|\bbrief\b|\bfocused\b|\bsingle\s+complaint\b",
    re.I,
)
_ACUTE_ILLNESS_RE = re.compile(
    r"\bfever\b|\bcough\b|\bsore\s+throat\b|\botitis\b|\bsprain\b|\bURI\b|"
    r"\bviral\b|\bacute\b|\binfection\b|\bpain\b|\brash\b|\bankle\b",
    re.I,
)
_CONSULT_RE = re.compile(r"\bconsult|\breferral\b|\breferred\b", re.I)
_GENERAL_ASSESS_RE = re.compile(
    r"\bgeneral\s+assessment\b|\bcomplete\s+assessment\b|\bcomprehensive\b|"
    r"\bnew\s+(?:adult\s+)?patient\b",
    re.I,
)
_REASSESS_RE = re.compile(
    r"\bre-?assessment\b|\breassess|\bgeneral\s+re\b",
    re.I,
)
_INTERMEDIATE_RE = re.compile(r"\bintermediate\b", re.I)
_MINI_RE = re.compile(r"\bmini\b", re.I)

_CONSULT_CODES = frozenset({"A006", "C005", "C006", "W105", "W106"})
_GENERAL_ASSESS_CODES = frozenset({"A003", "C003"})
_REASSESS_CODES = frozenset({"A004", "C004", "W004"})
_PERIODIC_CODES = frozenset({"K017", "W109"})
_MINOR_CODES = frozenset({"A001"})
_INTERMEDIATE_CODES = frozenset({"A007"})
_MINI_CODES = frozenset({"A008"})

# Core ambulatory ladders — keep special-visit / house-call codes from crowding
# top-5 when the note is a normal office visit. Consult codes are gated separately.
_OFFICE_LADDER_CODES = frozenset(
    {
        "A001",
        "A002",
        "A003",
        "A004",
        "A007",
        "A008",
        "K017",
        "A261",
        "A262",
        "A263",
        "A264",
        "A268",
        "A260",
        "A265",
        "A266",
        "A565",
        "A661",
    }
)
_OFFICE_ENCOUNTER_RE = re.compile(
    r"\bambulatory\b|\boffice\b|\boutpatient\b|\bclinic\b|\bminor\b|"
    r"\bintermediate\b|\bgeneral\b|\bperiodic\b|\bwell[\s-]?baby\b|"
    r"\blevel\s*[12]\b|\bmini\b",
    re.I,
)


def select_primary_docs(
    context: list[dict],
    *,
    learned_priors: dict[str, float] | None = None,
    limit: int | None = None,
    is_primary: PrimaryPred | None = None,
    clinical_summary: str | None = None,
    encounter_type: str | None = None,
    age_months: int | None = None,
    prefixes: list[str] | None = None,
) -> list[dict]:
    """Pick top primary codes from ranked retrieval (+ deterministic prior boost).

    Evidence gates bind codes to the encounter type, note language, age, and
    setting prefixes — same tightness goal as Paediatrics specialty ranking.
    """
    limit = limit or settings.deterministic_primary_count
    priors = learned_priors or {}
    primary_ok = is_primary or is_paediatrics_26_primary
    note = f"{encounter_type or ''} {clinical_summary or ''}"
    has_18m = bool(_ENHANCED_18M_RE.search(note))
    has_well_baby = bool(_WELL_BABY_RE.search(note))
    has_periodic = bool(_PERIODIC_RE.search(note))
    has_ltc = bool(_LTC_RE.search(note))
    has_minor = bool(_MINOR_HINT_RE.search(note))
    has_acute = bool(_ACUTE_ILLNESS_RE.search(clinical_summary or ""))
    has_consult = bool(_CONSULT_RE.search(note))
    has_general = bool(_GENERAL_ASSESS_RE.search(note))
    has_reassess = bool(_REASSESS_RE.search(note))
    has_intermediate = bool(_INTERMEDIATE_RE.search(note))
    has_mini = bool(_MINI_RE.search(note))
    age_in_18m_window = age_months is not None and 15 <= age_months <= 21
    allowed_prefixes = {p.upper()[:1] for p in (prefixes or []) if p}

    scored: list[tuple[float, str, dict]] = []
    for idx, doc in enumerate(context):
        if not is_current_sob(doc) or not primary_ok(doc):
            continue
        code = (doc.get("billing_code") or "").upper()
        if not code:
            continue
        # Re-assert encounter setting (A office / C hospital / W LTC / K …).
        if allowed_prefixes and code[:1] not in allowed_prefixes:
            continue

        # Lower index = stronger retrieval; priors add a stable boost.
        score = 1000.0 - float(idx) + float(priors.get(code, 0.0)) * 100.0

        # --- Hard evidence gates (drop when elements are unmet) ---

        # Enhanced 18-month codes: require note/encounter evidence (or tight age).
        if code in _ENHANCED_18M_CODES:
            if has_18m:
                score += 900.0
            elif age_in_18m_window and has_well_baby:
                score += 400.0
            else:
                continue

        # Consult / referral listings only when consult language is present.
        if code in _CONSULT_CODES and not has_consult:
            continue

        # Periodic health only with periodic/complete-physical language.
        if code in _PERIODIC_CODES:
            if not has_periodic:
                continue
            if code.startswith("W") and not has_ltc:
                continue
            if code == "K017" and has_ltc and not has_periodic:
                continue
            if code == "K017":
                score += 900.0

        # Office periodic → keep LTC W* out even if somehow prefix-allowed.
        if has_periodic and not has_ltc and code.startswith("W"):
            continue

        # General assessment / re-assessment require matching encounter cues.
        if code in _GENERAL_ASSESS_CODES and not (has_general or has_periodic):
            continue
        if code in _REASSESS_CODES and not has_reassess:
            continue

        # Minor / brief / single-complaint: hard-drop heavy ladder + consults.
        if has_minor and not has_general and not has_reassess and not has_consult:
            if code in (
                _GENERAL_ASSESS_CODES
                | _REASSESS_CODES
                | _CONSULT_CODES
                | _PERIODIC_CODES
                | _ENHANCED_18M_CODES
            ):
                continue
            if code in _MINOR_CODES:
                score += 500.0
            elif code in _MINI_CODES:
                score += 200.0
            elif code in _INTERMEDIATE_CODES and has_acute:
                score -= 120.0

        # Intermediate (without general/consult/periodic): prefer A007.
        if (
            has_intermediate
            and not has_general
            and not has_reassess
            and not has_consult
            and not has_periodic
            and not has_18m
        ):
            if code in _GENERAL_ASSESS_CODES | _REASSESS_CODES | _CONSULT_CODES:
                continue
            if code in _INTERMEDIATE_CODES:
                score += 450.0
            if has_acute and not has_well_baby and code in _ENHANCED_18M_CODES:
                continue

        # Mini assessment cue.
        if has_mini and code in _MINI_CODES:
            score += 450.0

        # General assessment cue.
        if has_general and code in _GENERAL_ASSESS_CODES:
            score += 500.0

        # Re-assessment cue.
        if has_reassess and code in _REASSESS_CODES:
            score += 500.0

        # Consult cue.
        if has_consult and code in _CONSULT_CODES:
            score += 500.0

        # Prefer schedule assessment ladder over special-visit oddities.
        if _OFFICE_ENCOUNTER_RE.search(note):
            if code in _OFFICE_LADDER_CODES or (
                has_consult and code in _CONSULT_CODES
            ):
                score += 120.0
            elif code[:1] == "A":
                score -= 250.0

        scored.append((score, code, doc))
    scored.sort(key=lambda t: (-t[0], t[1]))
    return [doc for _, _, doc in scored[:limit]]


def select_premium_docs(
    premium_pool: list[dict],
    *,
    primary_docs: list[dict],
    clinical_summary: str,
    time_align: Any | None,
    age_align: AgeAlignment | None,
    limit: int | None = None,
    is_premium: PremiumPred | None = None,
) -> list[dict]:
    """Rule-scored premiums/add-ons with stable tie-break on billing_code."""
    limit = limit or settings.deterministic_premium_count
    premium_ok = is_premium or is_paediatrics_26_premium
    note = clinical_summary or ""
    chronic = bool(_CHRONIC_RE.search(note))
    immuniz = bool(_IMMUNIZ_RE.search(note))
    procedure = bool(_PROCEDURE_RE.search(note))
    discharge = bool(_DISCHARGE_RE.search(note))

    related: list[str] = []
    for doc in primary_docs[:3]:
        for code in doc.get("relationship_premiums") or []:
            if code not in related:
                related.append(code)
    related_rank = {c: i for i, c in enumerate(related)}

    by_code: dict[str, dict] = {}
    for idx, doc in enumerate(premium_pool):
        if not is_current_sob(doc) or not premium_ok(doc):
            continue
        code = (doc.get("billing_code") or "").upper()
        if not code:
            continue
        # Evidence gates — stop relationship bleed (e.g. E078 on acute otitis).
        if code in _CHRONIC_CODES and not chronic:
            continue
        if code in _IMMUNIZ_CODES and not immuniz:
            continue
        if code in _DISCHARGE_CODES and not discharge:
            continue
        if code in _PROCEDURE_PREMIUMS and not procedure:
            continue
        if code in _AGE_ADDON_CODES and not age_align:
            continue
        if time_align and not premium_applies_to_band(doc, time_align.band):
            continue
        if age_align and not premium_applies_to_age(doc, age_align.age_months):
            continue
        # Prefer first occurrence (retrieval order).
        by_code.setdefault(code, {**doc, "_retrieval_idx": idx})

    scored: list[tuple[float, str, dict]] = []
    for code, doc in by_code.items():
        score = 0.0
        idx = int(doc.get("_retrieval_idx") or 0)
        score += max(0.0, 40.0 - float(idx))

        bands = premium_time_bands(doc)
        if time_align and bands and time_align.band in bands:
            score += 1000.0
        elif bands and _TIME_PREMIUM_HINT.search(
            f"{doc.get('description_text', '')} {doc.get('parent_section', '')}"
        ):
            # Time-banded but not matching — already filtered; skip residual.
            continue

        # Age-windowed premiums that survived age filter.
        if age_align and _looks_age_banded(doc):
            score += 500.0

        if code in _CHRONIC_CODES and chronic:
            score += 400.0
        if code in _IMMUNIZ_CODES and immuniz:
            score += 400.0
        if code in _DISCHARGE_CODES and discharge:
            score += 400.0
        if code in _PROCEDURE_PREMIUMS and procedure:
            score += 350.0

        if code in related_rank:
            score += 200.0 - float(related_rank[code])

        if score <= 0:
            continue
        scored.append((score, code, doc))

    scored.sort(key=lambda t: (-t[0], t[1]))
    return [doc for _, _, doc in scored[:limit]]


def build_deterministic_response(
    primary_docs: list[dict],
    premium_docs: list[dict],
    *,
    clinical_summary: str,
    time_align: Any | None = None,
    age_align: AgeAlignment | None = None,
    specialty_label: str = "Paediatrics",
    encounter_type: str | None = None,
) -> CaseAnalysisResponse:
    """Skeleton response with Schedule-grounded template justifications."""
    primaries: list[CodeRecommendation] = []
    for i, doc in enumerate(primary_docs):
        code = doc["billing_code"]
        desc = (doc.get("description_text") or code).strip()
        fee = float(doc.get("base_fee_cad") or 0.0)
        ci = _PRIMARY_CI[min(i, len(_PRIMARY_CI) - 1)]
        primaries.append(
            CodeRecommendation(
                rank=i + 1,
                billing_code=code,
                description=desc,
                base_fee_cad=fee,
                confidence_interval=ci,
                justification=_template_primary_justification(
                    doc,
                    clinical_summary,
                    rank=i + 1,
                    specialty_label=specialty_label,
                    encounter_type=encounter_type,
                ),
                reference=doc.get("reference"),
                differentiators=doc.get("differentiators"),
                parent_section=doc.get("parent_section"),
                fee_components=doc.get("fee_components"),
                rules_and_constraints=sob_grounded_rules(doc),
                in_current_schedule=True,
            )
        )

    associated: list[AssociatedCode] = []
    for i, doc in enumerate(premium_docs):
        code = doc["billing_code"]
        desc = (doc.get("description_text") or code).strip()
        rel = _relationship_label(doc)
        associated.append(
            AssociatedCode(
                billing_code=code,
                description=desc,
                relationship=rel,
                confidence_interval=_PREMIUM_CI[min(i, len(_PREMIUM_CI) - 1)],
                justification=_template_premium_justification(
                    doc,
                    clinical_summary,
                    time_align=time_align,
                    age_align=age_align,
                    encounter_type=encounter_type,
                ),
                base_fee_cad=doc.get("base_fee_cad"),
                reference=doc.get("reference"),
                differentiators=doc.get("differentiators"),
                parent_section=doc.get("parent_section"),
                fee_components=doc.get("fee_components"),
                rules_and_constraints=sob_grounded_rules(doc),
                in_current_schedule=True,
            )
        )

    return CaseAnalysisResponse(
        top_matching_codes=primaries,
        associated_codes=associated,
    )


def merge_justifications(
    result: CaseAnalysisResponse,
    justifications: dict[str, str],
) -> CaseAnalysisResponse:
    """Overlay LLM justifications onto a fixed code set (codes unchanged)."""
    for rec in result.top_matching_codes:
        text = (justifications.get(rec.billing_code) or "").strip()
        if text:
            rec.justification = text
    for assoc in result.associated_codes:
        text = (justifications.get(assoc.billing_code) or "").strip()
        if text:
            assoc.justification = text
    return result


def _looks_age_banded(doc: dict) -> bool:
    blob = " ".join(
        str(doc.get(k) or "")
        for k in ("description_text", "rules_and_constraints", "parent_section")
    )
    return bool(
        re.search(
            r"\bage|\byears?\b|\bmonths?\b|\binfant\b|\bpaediatric\b|\bpediatric\b",
            blob,
            re.I,
        )
    )


def _relationship_label(doc: dict) -> str:
    section = (doc.get("parent_section") or "").upper()
    code = (doc.get("billing_code") or "").upper()
    if code in _IMMUNIZ_CODES or "IMMUNIZATION" in section or "INJECTION" in section:
        return "add-on"
    return "premium"


_STOPWORDS = frozenset(
    {
        "a",
        "an",
        "the",
        "and",
        "or",
        "of",
        "to",
        "for",
        "in",
        "on",
        "with",
        "this",
        "that",
        "is",
        "are",
        "was",
        "were",
        "be",
        "by",
        "as",
        "at",
        "from",
        "into",
        "patient",
        "visit",
        "office",
        "care",
        "service",
        "assessment",
        "years",
        "year",
        "months",
        "month",
        "old",
        "per",
        "unit",
    }
)
_TOKEN_RE = re.compile(r"[A-Za-z][A-Za-z0-9'/-]{2,}")


def _note_evidence(summary: str, doc: dict, *, limit: int = 4) -> list[str]:
    """Content words from the note that also appear in the code's SoB text."""
    note = (summary or "").lower()
    if not note.strip():
        return []
    blob = " ".join(
        str(doc.get(k) or "")
        for k in (
            "description_text",
            "rules_and_constraints",
            "parent_section",
            "relationship_ladder",
            "relationship_family",
        )
    ).lower()
    note_tokens = [
        t for t in _TOKEN_RE.findall(note) if t not in _STOPWORDS and not t.isdigit()
    ]
    # Prefer longer clinical phrases first (2–3 grams) that appear in both.
    cues: list[str] = []
    words = note.split()
    for n in (3, 2):
        for i in range(max(0, len(words) - n + 1)):
            phrase = " ".join(words[i : i + n]).strip(".,;:()[]\"'")
            key = phrase.lower()
            if len(key) < 8:
                continue
            if key in blob and phrase not in cues:
                cues.append(phrase)
            if len(cues) >= limit:
                return cues
    for tok in note_tokens:
        if tok in blob and tok not in {c.lower() for c in cues}:
            cues.append(tok)
        if len(cues) >= limit:
            break
    return cues


def _diff_bits(doc: dict) -> list[str]:
    diff = doc.get("differentiators") or {}
    if not isinstance(diff, dict):
        return []
    out: list[str] = []
    if diff.get("time_band"):
        out.append(f"time band {diff['time_band']}")
    if diff.get("person_seen"):
        out.append(str(diff["person_seen"]))
    if diff.get("cap"):
        out.append(f"cap {diff['cap']}")
    return out


def _constraint_snippet(doc: dict, *, max_len: int = 160) -> str | None:
    rules = sob_grounded_rules(doc)
    if not rules:
        return None
    # Prefer the first sentence-like clause.
    text = re.split(r"(?<=[.|;])\s+", rules.strip())[0].strip()
    text = re.sub(r"\s+", " ", text)
    if len(text) > max_len:
        text = text[: max_len - 1].rsplit(" ", 1)[0] + "…"
    return text or None


def _template_primary_justification(
    doc: dict,
    summary: str,
    *,
    rank: int,
    specialty_label: str = "Paediatrics",
    encounter_type: str | None = None,
) -> str:
    """Clinician-facing explanation from SoB fields + note cues (no LLM)."""
    code = (doc.get("billing_code") or "").upper()
    desc = (doc.get("description_text") or code).strip()
    section = (doc.get("parent_section") or "").strip()
    ref = (doc.get("reference") or "").strip()
    fee = doc.get("base_fee_cad")
    sentences: list[str] = []

    head = f"{code} — {desc}."
    if fee is not None:
        try:
            head = f"{code} — {desc} (${float(fee):.2f})."
        except (TypeError, ValueError):
            pass
    sentences.append(head)

    why_bits: list[str] = []
    if encounter_type:
        why_bits.append(f"fits the selected encounter ({encounter_type})")
    cues = _note_evidence(summary, doc)
    if cues:
        why_bits.append("Schedule text aligns with note cues “" + "”, “".join(cues[:3]) + "”")
    snippet = " ".join((summary or "").split())
    if snippet:
        if len(snippet) > 110:
            snippet = snippet[:109].rsplit(" ", 1)[0] + "…"
        why_bits.append(f"clinical note: “{snippet}”")
    family = (doc.get("relationship_family") or "").replace("_", " ").strip()
    if family:
        why_bits.append(f"maps to the {family} service family for {specialty_label}")
    if why_bits:
        sentences.append("Recommended because " + "; ".join(why_bits) + ".")

    ladder = (doc.get("relationship_ladder") or "").strip()
    if ladder:
        sentences.append(ladder.rstrip(".") + ".")

    diffs = _diff_bits(doc)
    constraint = _constraint_snippet(doc)
    schedule_bits: list[str] = []
    if section:
        schedule_bits.append(f"SoB section “{section}”")
    if ref:
        schedule_bits.append(f"ref {ref}")
    if diffs:
        schedule_bits.append("differentiators: " + "; ".join(diffs))
    if schedule_bits:
        sentences.append("Grounded in " + ", ".join(schedule_bits) + ".")
    if constraint and constraint.lower() not in " ".join(sentences).lower():
        sentences.append(f"Constraint: {constraint}")

    if rank > 1:
        sentences.append(
            f"Listed as alternative #{rank} if a higher-ranked code’s elements "
            "are not fully met."
        )
    return " ".join(sentences)


def _template_premium_justification(
    doc: dict,
    summary: str,
    *,
    time_align: Any | None,
    age_align: AgeAlignment | None,
    encounter_type: str | None = None,
) -> str:
    """Clinician-facing premium/add-on explanation from rules + visit evidence."""
    code = (doc.get("billing_code") or "").upper()
    desc = (doc.get("description_text") or code).strip()
    rel = _relationship_label(doc)
    sentences: list[str] = [
        f"{code} — {desc} ({rel} alongside the primary visit code)."
    ]

    evidence: list[str] = []
    note = summary or ""
    if code in _IMMUNIZ_CODES and _IMMUNIZ_RE.search(note):
        evidence.append("the note documents immunization/injection")
    if code in _CHRONIC_CODES and _CHRONIC_RE.search(note):
        evidence.append("the note documents chronic disease follow-up")
    if code in _DISCHARGE_CODES and _DISCHARGE_RE.search(note):
        evidence.append("the note documents post-hospital/discharge context")
    if code in _PROCEDURE_PREMIUMS and _PROCEDURE_RE.search(note):
        evidence.append("the note documents a procedure")
    if time_align and premium_time_bands(doc) and time_align.band in premium_time_bands(doc):
        evidence.append(f"Time of Service matches {time_align.label}")
    if age_align and (
        code in _AGE_ADDON_CODES or _looks_age_banded(doc)
    ):
        evidence.append(f"patient age ({age_align.label}) is in the fee’s age window")
    cues = _note_evidence(note, doc, limit=2)
    if cues and not evidence:
        evidence.append("note cues “" + "”, “".join(cues) + "”")
    if encounter_type and not evidence:
        evidence.append(f"compatible with {encounter_type}")
    if evidence:
        sentences.append("Eligible because " + "; ".join(evidence) + ".")

    constraint = _constraint_snippet(doc, max_len=140)
    if constraint:
        sentences.append(f"Constraint: {constraint}")
    ref = (doc.get("reference") or "").strip()
    section = (doc.get("parent_section") or "").strip()
    if ref or section:
        loc = []
        if section:
            loc.append(section)
        if ref:
            loc.append(f"ref {ref}")
        sentences.append("See " + " · ".join(loc) + ".")
    return " ".join(sentences)