decision/engine/domain_nomination.py

"""
Domain Nomination Layer.

Combines DriveHealthBERT ML output with taxonomy trigger matches to produce
a ranked list of clinical domain nominations.

The nomination algorithm:
  1. Run taxonomy triggers against the text → domain match results
  2. Receive DriveHealthBERT label + probabilities
  3. Fuse signals: ML confidence reinforces or weakens trigger confidence
  4. Apply hard-escalate rules (suicidal_ideation, homicidal_ideation)
  5. Rank by: (a) effective_confidence tier, (b) domain priority, (c) ML prob
  6. Select recommended_flow from domain rules based on confidence tier
  7. Return sorted list of DomainNomination objects

Safety invariants:
  - Hard-escalate domains ALWAYS nominate at HIGH confidence
  - Negated domains are included but marked (caller decides)
  - If ML says ESCALATION but no triggers match, still nominate with
    a generic "unclassified_escalation" domain
  - If triggers match a safety domain but ML says non-escalation,
    the trigger OVERRIDES the ML (fail-open)
"""

from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional, Set, Tuple

from decision.engine.config_loader import DecisionConfigLoader
from decision.engine.models import (
    ConfidenceTier,
    DomainNomination,
    NegationResult,
    RiskClass,
    TriggerMatch,
)
from decision.engine.trigger_engine import DomainMatchResult, TaxonomyTriggerEngine

logger = logging.getLogger("decision.domain_nomination")

# DriveHealthBERT label → taxonomy domain mapping hints
# Used when ML label alone doesn't resolve to a specific domain
_ML_LABEL_DOMAIN_HINTS: Dict[str, List[str]] = {
    "SCHEDULING": ["scheduling", "scheduling_barrier", "missed_followup"],
    "MEDICATION": ["medication", "medication_nonadherence"],
    "SYMPTOM_CHECK": [],  # Too broad — rely on triggers
    "BILLING": [],         # No taxonomy domain for billing yet
    "GENERAL": [],
    "ESCALATION": [],      # Rely on triggers to determine which domain
}


class DomainNominator:
    """
    Produces ranked domain nominations from ML + lexical signals.

    Usage:
        nominator = DomainNominator(config, trigger_engine)
        nominations = nominator.nominate(
            text="my chest is killing me",
            ml_label="ESCALATION",
            ml_probabilities={"ESCALATION": 0.92, ...},
        )
    """

    def __init__(
        self,
        config: DecisionConfigLoader,
        trigger_engine: TaxonomyTriggerEngine,
    ):
        self._config = config
        self._trigger_engine = trigger_engine
        self._hard_escalate: Set[str] = config.hard_escalate_domains
        self._safety_precedence: List[str] = config.safety_precedence

    def nominate(
        self,
        text: str,
        ml_label: Optional[str] = None,
        ml_probabilities: Optional[Dict[str, float]] = None,
    ) -> List[DomainNomination]:
        """
        Produce ranked domain nominations for the given text.

        Args:
            text: Patient utterance (original, NOT context-prepended)
            ml_label: DriveHealthBERT predicted label (e.g., "ESCALATION")
            ml_probabilities: Full probability distribution from DriveHealthBERT

        Returns:
            List of DomainNomination sorted by:
              1. Effective confidence tier (HIGH > MEDIUM > LOW)
              2. Domain priority (higher = more urgent)
              3. ML probability (if available)
        """
        ml_confidence = None
        if ml_probabilities and ml_label:
            ml_confidence = ml_probabilities.get(ml_label, 0.0)

        # Step 1: Run taxonomy triggers
        trigger_results = self._trigger_engine.match_all(text)

        # Step 2: Build nominations from trigger results
        nominations: List[DomainNomination] = []

        for domain_name, match_result in trigger_results.items():
            nomination = self._build_nomination(
                domain_name=domain_name,
                match_result=match_result,
                ml_label=ml_label,
                ml_confidence=ml_confidence,
            )
            nominations.append(nomination)

        # Step 3: Handle hard-escalate domains
        #   If any hard-escalate domain matched (even at LOW), force HIGH
        for nom in nominations:
            if nom.domain in self._hard_escalate and not nom.is_negated:
                nominations = [
                    self._force_high_confidence(n) if n.domain == nom.domain else n
                    for n in nominations
                ]

        # Step 4: Handle ML-only signals (ML says ESCALATION but no triggers)
        if ml_label == "ESCALATION" and ml_confidence and ml_confidence > 0.5:
            has_safety_trigger = any(
                not n.is_negated and n.confidence_tier >= ConfidenceTier.LOW
                for n in nominations
                if n.domain in set(self._safety_precedence)
            )
            # Check if the ML is confused by negated safety keywords:
            # If safety domains DID match but were ALL negated, the ML is
            # likely reacting to the keywords themselves (e.g., "I'm NOT
            # suicidal" triggers ML ESCALATION because it sees "suicidal").
            # In this case, trust the explicit negation over the ML signal.
            has_negated_safety = any(
                n.is_negated
                for n in nominations
                if n.domain in set(self._safety_precedence)
            )
            if not has_safety_trigger and not has_negated_safety:
                # ML detected escalation but no specific domain matched
                # (and no negated safety domains either)
                # Create a generic nomination so it's not silently dropped
                nominations.append(
                    DomainNomination(
                        domain="unclassified_escalation",
                        display_name="Unclassified Escalation",
                        confidence_tier=self._ml_prob_to_tier(ml_confidence),
                        priority=50,  # Middle priority
                        target_risk_class=RiskClass.R2,
                        trigger_matches=(),
                        negation_result=None,
                        ml_label=ml_label,
                        ml_confidence=ml_confidence,
                        recommended_flow="handoff",
                    )
                )

        # Step 5: Handle ML reinforcement for non-escalation labels
        if ml_label and ml_label != "ESCALATION":
            hint_domains = _ML_LABEL_DOMAIN_HINTS.get(ml_label, [])
            for nom in nominations:
                if nom.domain in hint_domains and ml_confidence:
                    # ML confirms trigger match — boost confidence if ML is strong
                    if ml_confidence > 0.8 and nom.confidence_tier < ConfidenceTier.HIGH:
                        nominations = [
                            self._boost_confidence(n) if n.domain == nom.domain else n
                            for n in nominations
                        ]

        # Step 6: Resolve recommended_flow for each nomination
        nominations = [self._resolve_flow(n) for n in nominations]

        # Step 7: Sort
        nominations.sort(key=self._nomination_sort_key, reverse=True)

        if nominations:
            logger.info(
                "Domain nominations for text (len=%d): %s",
                len(text),
                [(n.domain, n.confidence_tier.value, n.priority) for n in nominations[:5]],
            )

        return nominations

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _build_nomination(
        self,
        domain_name: str,
        match_result: DomainMatchResult,
        ml_label: Optional[str],
        ml_confidence: Optional[float],
    ) -> DomainNomination:
        """Build a DomainNomination from a DomainMatchResult."""
        # Get target risk class from rules.yaml
        target_risk_str = self._config.get_domain_target_risk(domain_name)
        target_risk = None
        if target_risk_str:
            try:
                target_risk = RiskClass(target_risk_str)
            except ValueError:
                pass

        return DomainNomination(
            domain=domain_name,
            display_name=domain_name.replace("_", " ").title(),
            confidence_tier=match_result.effective_confidence,
            priority=match_result.priority,
            target_risk_class=target_risk,
            trigger_matches=match_result.matches,
            negation_result=match_result.negation_result,
            ml_label=ml_label,
            ml_confidence=ml_confidence,
        )

    def _resolve_flow(self, nomination: DomainNomination) -> DomainNomination:
        """Determine recommended_flow from domain decision rules."""
        if nomination.recommended_flow:
            return nomination

        rules = self._config.get_domain_decision_rules(nomination.domain)
        if not rules:
            return nomination

        # Find first matching rule based on confidence tier
        for rule in rules:
            conditions = rule.get("if", {})
            conf_required = conditions.get("confidence")
            if conf_required and conf_required == nomination.confidence_tier.value:
                flow = rule.get("then", {}).get("flow")
                if flow:
                    return DomainNomination(
                        domain=nomination.domain,
                        display_name=nomination.display_name,
                        confidence_tier=nomination.confidence_tier,
                        priority=nomination.priority,
                        target_risk_class=nomination.target_risk_class,
                        trigger_matches=nomination.trigger_matches,
                        negation_result=nomination.negation_result,
                        ml_label=nomination.ml_label,
                        ml_confidence=nomination.ml_confidence,
                        recommended_flow=flow,
                    )

        return nomination

    def _force_high_confidence(
        self, nomination: DomainNomination
    ) -> DomainNomination:
        """Force a nomination to HIGH confidence (for hard-escalate domains)."""
        if nomination.confidence_tier == ConfidenceTier.HIGH:
            return nomination
        logger.warning(
            "HARD ESCALATE: Forcing %s to HIGH confidence (was %s)",
            nomination.domain,
            nomination.confidence_tier.value,
        )
        return DomainNomination(
            domain=nomination.domain,
            display_name=nomination.display_name,
            confidence_tier=ConfidenceTier.HIGH,
            priority=nomination.priority,
            target_risk_class=nomination.target_risk_class,
            trigger_matches=nomination.trigger_matches,
            negation_result=nomination.negation_result,
            ml_label=nomination.ml_label,
            ml_confidence=nomination.ml_confidence,
            recommended_flow=nomination.recommended_flow,
        )

    def _boost_confidence(
        self, nomination: DomainNomination
    ) -> DomainNomination:
        """Boost confidence by one tier when ML reinforces trigger."""
        new_tier = nomination.confidence_tier
        if new_tier == ConfidenceTier.LOW:
            new_tier = ConfidenceTier.MEDIUM
        elif new_tier == ConfidenceTier.MEDIUM:
            new_tier = ConfidenceTier.HIGH
        if new_tier == nomination.confidence_tier:
            return nomination
        return DomainNomination(
            domain=nomination.domain,
            display_name=nomination.display_name,
            confidence_tier=new_tier,
            priority=nomination.priority,
            target_risk_class=nomination.target_risk_class,
            trigger_matches=nomination.trigger_matches,
            negation_result=nomination.negation_result,
            ml_label=nomination.ml_label,
            ml_confidence=nomination.ml_confidence,
            recommended_flow=nomination.recommended_flow,
        )

    @staticmethod
    def _ml_prob_to_tier(prob: float) -> ConfidenceTier:
        """Map ML probability to a confidence tier."""
        if prob >= 0.8:
            return ConfidenceTier.HIGH
        elif prob >= 0.5:
            return ConfidenceTier.MEDIUM
        elif prob >= 0.3:
            return ConfidenceTier.LOW
        return ConfidenceTier.NONE

    @staticmethod
    def _nomination_sort_key(
        nom: DomainNomination,
    ) -> Tuple[int, int, float]:
        """Sort key: (confidence_rank, priority, ml_confidence)."""
        conf_rank = {
            ConfidenceTier.HIGH: 3,
            ConfidenceTier.MEDIUM: 2,
            ConfidenceTier.LOW: 1,
            ConfidenceTier.NONE: 0,
        }
        return (
            conf_rank.get(nom.confidence_tier, 0),
            nom.priority,
            nom.ml_confidence or 0.0,
        )