src/core/chaplain_models.py

# chaplain_models.py
"""
Data models for Chaplain Feedback & Tagging System.

Defines core data structures for classification flows, tagging records,
distress indicators, and interaction logging.
"""

from dataclasses import dataclass, field
from typing import List, Optional, Dict, Any
from datetime import datetime


# =============================================================================
# INDICATOR DEFINITIONS - Based on Spiritual Distress Definitions Document
# =============================================================================

# Mapping of all indicators from the definitions document with their categories,
# subcategories, severity (red/yellow), and definition references.
# RED (#ea9999): Severe distress - requires immediate attention
# YELLOW (#ffe599): Potential distress - requires clarification

INDICATOR_DEFINITIONS: Dict[str, Dict[str, Any]] = {
    # Section II.A - Emotional expressions
    "crying": {
        "category": "Emotional",
        "subcategory": "Crying",
        "severity": "red",
        "definition_reference": "II.A",
        "description": "Crying as expression of spiritual distress"
    },
    "dysomnias": {
        "category": "Emotional",
        "subcategory": "Dysomnias/Difficulty sleeping",
        "severity": "yellow",
        "definition_reference": "II.A",
        "description": "Sleep disturbances related to spiritual distress"
    },
    "fatigue": {
        "category": "Emotional",
        "subcategory": "Fatigue, emotional exhaustion",
        "severity": "yellow",
        "definition_reference": "II.A",
        "description": "Fatigue and emotional exhaustion"
    },
    "anxiety": {
        "category": "Emotional",
        "subcategory": "Anxiety",
        "severity": "yellow",
        "definition_reference": "II.A",
        "description": "Anxiety as expression of spiritual distress"
    },
    "fear": {
        "category": "Emotional",
        "subcategory": "Fear",
        "severity": "yellow",
        "definition_reference": "II.A",
        "description": "Fear as expression of spiritual distress"
    },
    "anger": {
        "category": "Emotional",
        "subcategory": "Anger",
        "severity": "red",
        "definition_reference": "II.A",
        "description": "Anger as expression of spiritual distress"
    },
    "depressive_symptoms": {
        "category": "Emotional",
        "subcategory": "Depressive symptoms",
        "severity": "yellow",
        "definition_reference": "II.A",
        "description": "Depressive symptoms"
    },
    
    # Section II.B - Decreased engagement
    "decreased_engagement": {
        "category": "Engagement",
        "subcategory": "Decreased engagement with hobbies",
        "severity": "yellow",
        "definition_reference": "II.B",
        "description": "Decreased engagement with hobbies, creative expression, and personal interests"
    },
    
    # Section II.C - Disinterest in nature
    "disinterest_nature": {
        "category": "Engagement",
        "subcategory": "Disinterest in nature",
        "severity": "yellow",
        "definition_reference": "II.C",
        "description": "Disinterest in nature due to spiritual, emotional and physical limitations"
    },
    
    # Section II.D - Excessive guilt
    "excessive_guilt": {
        "category": "Guilt",
        "subcategory": "Excessive guilt",
        "severity": "red",
        "definition_reference": "II.D",
        "description": "Excessive guilt - existential, religious, or relational"
    },
    
    # Section II.E - Anger behaviors of spiritual nature
    "anger_spiritual": {
        "category": "Anger",
        "subcategory": "Anger behaviors of a spiritual nature",
        "severity": "red",
        "definition_reference": "II.E",
        "description": "Anger toward power greater than self"
    },
    
    # Section II.F - Grief types
    "anticipatory_grieving": {
        "category": "Grief",
        "subcategory": "Anticipatory grieving",
        "severity": "red",
        "definition_reference": "II.F",
        "description": "Emotional response to anticipated death"
    },
    "disenfranchised_grief": {
        "category": "Grief",
        "subcategory": "Disenfranchised grief",
        "severity": "red",
        "definition_reference": "II.F",
        "description": "Grief unacknowledged or unsupported by society"
    },
    "life_review_grieving": {
        "category": "Grief",
        "subcategory": "Grieving in the setting of life review",
        "severity": "yellow",
        "definition_reference": "II.F",
        "description": "Grieving during life review process"
    },
    "maladaptive_grieving": {
        "category": "Grief",
        "subcategory": "Maladaptive grieving",
        "severity": "red",
        "definition_reference": "II.F",
        "description": "Prolonged grief disorder"
    },
    "complicated_grief": {
        "category": "Grief",
        "subcategory": "Complicated grief",
        "severity": "red",
        "definition_reference": "II.F",
        "description": "Persistent, intense grief disrupting daily life"
    },
    "loss_loved_one": {
        "category": "Grief",
        "subcategory": "Loss of a loved one",
        "severity": "red",
        "definition_reference": "II.F",
        "description": "Loss of family member or friend"
    },

    # Section II.G - Expressions of Spiritual Distress
    "expresses_alienation": {
        "category": "Expressions",
        "subcategory": "Expresses alienation",
        "severity": "yellow",
        "definition_reference": "II.G",
        "description": "Feeling separation, isolation, disconnection"
    },
    "concern_beliefs": {
        "category": "Expressions",
        "subcategory": "Expresses concern about beliefs",
        "severity": "yellow",
        "definition_reference": "II.G",
        "description": "Questions or struggles with spiritual/religious beliefs"
    },
    "concern_future": {
        "category": "Expressions",
        "subcategory": "Expresses concern about the future",
        "severity": "red",
        "definition_reference": "II.G",
        "description": "Anxious, fearful, or uncertain about what lies ahead"
    },
    "concern_values": {
        "category": "Expressions",
        "subcategory": "Expresses concern about values system",
        "severity": "yellow",
        "definition_reference": "II.G",
        "description": "Conflicted about moral or ethical principles"
    },
    "concern_family": {
        "category": "Expressions",
        "subcategory": "Expresses concerns about family",
        "severity": "yellow",
        "definition_reference": "II.G",
        "description": "Distressed about family well-being or relationships"
    },
    "feeling_emptiness": {
        "category": "Expressions",
        "subcategory": "Expresses feeling of emptiness",
        "severity": "red",
        "definition_reference": "II.G",
        "description": "Deep inner void or lack of meaning"
    },
    "feeling_unloved": {
        "category": "Expressions",
        "subcategory": "Expresses feeling unloved",
        "severity": "red",
        "definition_reference": "II.G",
        "description": "Feels unworthy of love or disconnected from caring relationships"
    },
    "feeling_worthless": {
        "category": "Expressions",
        "subcategory": "Expresses feeling worthless",
        "severity": "red",
        "definition_reference": "II.G",
        "description": "Perceives themselves as having little or no value"
    },
    "insufficient_courage": {
        "category": "Expressions",
        "subcategory": "Expresses insufficient courage",
        "severity": "yellow",
        "definition_reference": "II.G",
        "description": "Fear or lack of strength to face suffering"
    },
    "loss_confidence": {
        "category": "Expressions",
        "subcategory": "Expresses loss of confidence",
        "severity": "yellow",
        "definition_reference": "II.G",
        "description": "Diminished trust in themselves or abilities"
    },
    "loss_control": {
        "category": "Expressions",
        "subcategory": "Expresses loss of control",
        "severity": "yellow",
        "definition_reference": "II.G",
        "description": "Feels powerless over life circumstances"
    },
    "loss_hope": {
        "category": "Expressions",
        "subcategory": "Expresses loss of hope",
        "severity": "red",
        "definition_reference": "II.G",
        "description": "Feels despair or believes future holds no possibility"
    },
    "loss_serenity": {
        "category": "Expressions",
        "subcategory": "Expresses loss of serenity",
        "severity": "yellow",
        "definition_reference": "II.G",
        "description": "Inner turmoil, anxiety, or restlessness"
    },
    "need_forgiveness": {
        "category": "Expressions",
        "subcategory": "Expresses need for forgiveness",
        "severity": "red",
        "definition_reference": "II.G",
        "description": "Feels guilt or remorse and desires reconciliation"
    },
    "expresses_regret": {
        "category": "Expressions",
        "subcategory": "Expresses regret",
        "severity": "yellow",
        "definition_reference": "II.G",
        "description": "Sorrow over past actions or missed opportunities"
    },
    "expresses_suffering": {
        "category": "Expressions",
        "subcategory": "Expresses suffering",
        "severity": "red",
        "definition_reference": "II.G",
        "description": "Deep physical, emotional, or spiritual pain"
    },
    "concern_medical_treatment": {
        "category": "Medical",
        "subcategory": "Expresses concern about medical treatment",
        "severity": "red",
        "definition_reference": "II.G",
        "description": "Concern about treatment or medical team"
    },
    "unfinished_business": {
        "category": "Expressions",
        "subcategory": "Expresses feeling of having unfinished business",
        "severity": "red",
        "definition_reference": "II.G",
        "description": "Important matters remain unresolved"
    },
    "desire_share_spiritual": {
        "category": "Spiritual",
        "subcategory": "Expresses desire to share intense spiritual experiences",
        "severity": "yellow",
        "definition_reference": "II.G",
        "description": "Wants to share intense spiritual/religious experiences"
    },
    "inability_transcendence": {
        "category": "Spiritual",
        "subcategory": "Inability to experience transcendence",
        "severity": "red",
        "definition_reference": "II.G",
        "description": "Cannot experience supportive forces larger than oneself"
    },
    "impaired_introspection": {
        "category": "Spiritual",
        "subcategory": "Impaired ability for introspection",
        "severity": "yellow",
        "definition_reference": "II.G",
        "description": "Impaired ability for self-reflection"
    },

    # Section II.H - Existential questioning
    "questioning_identity": {
        "category": "Existential",
        "subcategory": "Questioning one's identity",
        "severity": "yellow",
        "definition_reference": "II.H",
        "description": "Confused about identity when illness takes away roles"
    },
    "questioning_meaning_life": {
        "category": "Existential",
        "subcategory": "Questioning the meaning of life",
        "severity": "red",
        "definition_reference": "II.H",
        "description": "Grapples with fundamental questions about existence"
    },
    "questioning_meaning_suffering": {
        "category": "Existential",
        "subcategory": "Questioning the meaning of suffering",
        "severity": "red",
        "definition_reference": "II.H",
        "description": "Struggles to understand if pain has purpose"
    },
    "questioning_dignity": {
        "category": "Existential",
        "subcategory": "Questioning one's own dignity",
        "severity": "red",
        "definition_reference": "II.H",
        "description": "Questions inherent worth and value as person"
    },
    
    # Section II.I - Social isolation
    "social_isolation": {
        "category": "Social",
        "subcategory": "Social isolation expressions",
        "severity": "yellow",
        "definition_reference": "II.I",
        "description": "Avoids interaction, estrangement, loneliness"
    },
    
    # Section II.J - Changes in spiritual/religious practices
    "altered_religious_ritual": {
        "category": "Spiritual",
        "subcategory": "Altered religious ritual",
        "severity": "yellow",
        "definition_reference": "II.J.a",
        "description": "Disruption to religious practices"
    },
    "altered_spiritual_practice": {
        "category": "Spiritual",
        "subcategory": "Altered spiritual practice",
        "severity": "yellow",
        "definition_reference": "II.J.b",
        "description": "Disruption to personal spiritual activities"
    },
    
    # Section II.K - Cultural conflict
    "cultural_conflict": {
        "category": "Cultural",
        "subcategory": "Cultural conflict",
        "severity": "yellow",
        "definition_reference": "II.K",
        "description": "Clash between cultural beliefs and healthcare culture"
    },
    
    # Section II.L - Sociocultural deprivation
    "sociocultural_deprivation": {
        "category": "Cultural",
        "subcategory": "Sociocultural deprivation",
        "severity": "yellow",
        "definition_reference": "II.L",
        "description": "Separated from cultural community"
    },
    
    # Section II.M - Difficulty accepting aging
    "difficulty_accepting_aging": {
        "category": "Aging",
        "subcategory": "Difficulty accepting aging",
        "severity": "yellow",
        "definition_reference": "II.M",
        "description": "Grief over lost abilities, resistance to mortality"
    },
    
    # Section II.N - Inadequate environmental control
    "inadequate_environmental_control": {
        "category": "Environment",
        "subcategory": "Inadequate environmental control",
        "severity": "yellow",
        "definition_reference": "II.N",
        "description": "Unable to shape surroundings for spiritual needs"
    },
    
    # Section II.O - Loss of independence
    "loss_independence": {
        "category": "Independence",
        "subcategory": "Loss of independence",
        "severity": "yellow",
        "definition_reference": "II.O",
        "description": "Dependency threatens personal and spiritual agency"
    },
    
    # Section II.P - Uncontrolled pain
    "uncontrolled_pain": {
        "category": "Medical",
        "subcategory": "Uncontrolled pain",
        "severity": "red",
        "definition_reference": "II.P",
        "description": "Persistent physical pain causing existential distress"
    },
    
    # Section II.Q - Spiritual pain
    "spiritual_pain": {
        "category": "Spiritual",
        "subcategory": "Spiritual pain",
        "severity": "red",
        "definition_reference": "II.Q",
        "description": "Soul-level suffering beyond physical symptoms"
    },
}


# =============================================================================
# DATA MODELS
# =============================================================================

@dataclass
class DistressIndicator:
    """
    Detected distress indicator with category and severity.
    
    Based on the Spiritual Distress Definitions document with color coding:
    - RED (#ea9999): Severe distress - requires immediate attention
    - YELLOW (#ffe599): Potential distress - requires clarification
    """
    indicator_text: str
    category: str  # "Emotional", "Grief", "Existential", "Expressions", "Spiritual", "Medical", "Social", "Cultural"
    subcategory: str  # Specific indicator name from definitions document
    severity: str  # "red" or "yellow" - based on color coding in definitions document
    confidence: float  # 0.0-1.0
    definition_reference: str = ""  # Section reference (e.g., "II.D", "II.G")
    
    def __post_init__(self):
        """Validate severity value."""
        if self.severity not in ("red", "yellow"):
            raise ValueError(f"Severity must be 'red' or 'yellow', got '{self.severity}'")
        if not 0.0 <= self.confidence <= 1.0:
            raise ValueError(f"Confidence must be between 0.0 and 1.0, got {self.confidence}")
    
    def to_dict(self) -> dict:
        """Convert indicator to dictionary for serialization."""
        return {
            "indicator_text": self.indicator_text,
            "category": self.category,
            "subcategory": self.subcategory,
            "severity": self.severity,
            "confidence": self.confidence,
            "definition_reference": self.definition_reference,
        }
    
    @classmethod
    def from_dict(cls, data: dict) -> "DistressIndicator":
        """Create indicator from dictionary."""
        return cls(**data)
    
    @classmethod
    def from_definition(cls, indicator_key: str, indicator_text: str, confidence: float) -> "DistressIndicator":
        """
        Create indicator from INDICATOR_DEFINITIONS constant.
        
        Args:
            indicator_key: Key in INDICATOR_DEFINITIONS (e.g., "excessive_guilt")
            indicator_text: The actual text that triggered this indicator
            confidence: Confidence score 0.0-1.0
        
        Returns:
            DistressIndicator with category, subcategory, severity from definitions
        
        Raises:
            KeyError: If indicator_key not found in INDICATOR_DEFINITIONS
        """
        if indicator_key not in INDICATOR_DEFINITIONS:
            raise KeyError(f"Unknown indicator key: {indicator_key}")
        
        defn = INDICATOR_DEFINITIONS[indicator_key]
        return cls(
            indicator_text=indicator_text,
            category=defn["category"],
            subcategory=defn["subcategory"],
            severity=defn["severity"],
            confidence=confidence,
            definition_reference=defn["definition_reference"],
        )



@dataclass
class FollowUpQuestion:
    """
    Generated follow-up question for YELLOW cases.
    
    Contains 1-2 short, sensitive clarifying questions with purpose explanation.
    """
    question_id: str
    question_text: str
    purpose: str  # Why this question is being asked
    
    def to_dict(self) -> dict:
        """Convert question to dictionary for serialization."""
        return {
            "question_id": self.question_id,
            "question_text": self.question_text,
            "purpose": self.purpose,
        }
    
    @classmethod
    def from_dict(cls, data: dict) -> "FollowUpQuestion":
        """Create question from dictionary."""
        return cls(**data)


@dataclass
class ClassificationFlowResult:
    """
    Complete result of classification flow.
    
    Contains all flow-specific fields for RED/YELLOW/GREEN classifications.
    """
    classification: str  # "red", "yellow", "green"
    confidence: float  # 0.0-1.0
    indicators: List[DistressIndicator] = field(default_factory=list)
    explanation: str = ""
    
    # RED-specific fields
    permission_check_message: Optional[str] = None
    referral_message: Optional[str] = None
    consent_status: Optional[str] = None  # "granted", "declined", None
    
    # YELLOW-specific fields
    follow_up_questions: List[FollowUpQuestion] = field(default_factory=list)
    patient_responses: List[str] = field(default_factory=list)
    re_evaluation_result: Optional[str] = None  # "red", "green", None
    
    def __post_init__(self):
        """Validate classification value."""
        if self.classification not in ("red", "yellow", "green"):
            raise ValueError(f"Classification must be 'red', 'yellow', or 'green', got '{self.classification}'")
        if not 0.0 <= self.confidence <= 1.0:
            raise ValueError(f"Confidence must be between 0.0 and 1.0, got {self.confidence}")
    
    def to_dict(self) -> dict:
        """Convert result to dictionary for serialization."""
        return {
            "classification": self.classification,
            "confidence": self.confidence,
            "indicators": [i.to_dict() for i in self.indicators],
            "explanation": self.explanation,
            "permission_check_message": self.permission_check_message,
            "referral_message": self.referral_message,
            "consent_status": self.consent_status,
            "follow_up_questions": [q.to_dict() for q in self.follow_up_questions],
            "patient_responses": self.patient_responses,
            "re_evaluation_result": self.re_evaluation_result,
        }
    
    @classmethod
    def from_dict(cls, data: dict) -> "ClassificationFlowResult":
        """Create result from dictionary."""
        data_copy = data.copy()
        
        # Convert nested indicators
        indicators_data = data_copy.pop("indicators", [])
        indicators = [DistressIndicator.from_dict(i) for i in indicators_data]
        
        # Convert nested follow-up questions
        questions_data = data_copy.pop("follow_up_questions", [])
        questions = [FollowUpQuestion.from_dict(q) for q in questions_data]
        
        result = cls(**data_copy)
        result.indicators = indicators
        result.follow_up_questions = questions
        return result


# Tagging category constants
CLASSIFICATION_SUBCATEGORIES = [
    "missed_indicators",  # Missed key distress indicators
    "false_positive",     # Overly sensitive (false-positive flag)
    "missed_distress",    # Not sensitive enough (missed distress)
]

QUESTION_ISSUE_TYPES = [
    "inappropriate",      # Question is inappropriate or intrusive
    "not_relevant",       # Question is not spiritually relevant
    "too_leading",        # Question is too leading or assumptive
    "unclear",            # Question is unclear or confusing
    "tone_clinical",      # Tone too clinical
    "tone_religious",     # Tone too religious
    "tone_casual",        # Tone too casual
]

REFERRAL_ISSUE_TYPES = [
    "incomplete_summary",    # Incorrect or incomplete summary
    "misrepresentation",     # Misrepresentation of patient message
    "inappropriate_tone",    # Tone inappropriate for spiritual care team
]


@dataclass
class TaggingRecord:
    """
    Structured tagging feedback from chaplain.
    
    Supports multi-select for question and referral issues.
    """
    record_id: str
    message_id: str
    
    # Classification feedback
    is_classification_correct: bool = True
    classification_subcategory: Optional[str] = None  # "missed_indicators", "false_positive", "missed_distress"
    correct_classification: Optional[str] = None  # "red", "yellow", "green"
    
    # Follow-up question feedback (YELLOW only)
    question_issues: List[str] = field(default_factory=list)  # Multi-select from QUESTION_ISSUE_TYPES
    question_comments: Optional[str] = None
    
    # Referral message feedback (RED only)
    referral_issues: List[str] = field(default_factory=list)  # Multi-select from REFERRAL_ISSUE_TYPES
    referral_comments: Optional[str] = None
    
    # Indicator feedback
    indicator_issues: List[str] = field(default_factory=list)  # List of incorrectly identified indicator IDs
    indicator_comments: Optional[str] = None
    
    # General
    general_notes: str = ""
    timestamp: datetime = field(default_factory=datetime.now)
    
    def __post_init__(self):
        """Validate tagging values."""
        if self.classification_subcategory and self.classification_subcategory not in CLASSIFICATION_SUBCATEGORIES:
            raise ValueError(f"Invalid classification subcategory: {self.classification_subcategory}")
        if self.correct_classification and self.correct_classification not in ("red", "yellow", "green"):
            raise ValueError(f"Invalid correct_classification: {self.correct_classification}")
        for issue in self.question_issues:
            if issue not in QUESTION_ISSUE_TYPES:
                raise ValueError(f"Invalid question issue type: {issue}")
        for issue in self.referral_issues:
            if issue not in REFERRAL_ISSUE_TYPES:
                raise ValueError(f"Invalid referral issue type: {issue}")
    
    def to_dict(self) -> dict:
        """Convert record to dictionary for serialization."""
        return {
            "record_id": self.record_id,
            "message_id": self.message_id,
            "is_classification_correct": self.is_classification_correct,
            "classification_subcategory": self.classification_subcategory,
            "correct_classification": self.correct_classification,
            "question_issues": self.question_issues,
            "question_comments": self.question_comments,
            "referral_issues": self.referral_issues,
            "referral_comments": self.referral_comments,
            "indicator_issues": self.indicator_issues,
            "indicator_comments": self.indicator_comments,
            "general_notes": self.general_notes,
            "timestamp": self.timestamp.isoformat(),
        }
    
    @classmethod
    def from_dict(cls, data: dict) -> "TaggingRecord":
        """Create record from dictionary."""
        data_copy = data.copy()
        if isinstance(data_copy.get("timestamp"), str):
            data_copy["timestamp"] = datetime.fromisoformat(data_copy["timestamp"])
        return cls(**data_copy)



# Interaction step types
INTERACTION_STEP_TYPES = [
    "classification",      # Initial classification
    "explanation",         # Explanation generation
    "permission_check",    # Patient consent request
    "follow_up",           # Follow-up questions
    "referral",            # Referral message generation
    "feedback",            # Chaplain feedback
]


@dataclass
class InteractionStepLog:
    """
    Log entry for a single interaction step.
    
    Records all interaction steps with input/output for analysis.
    """
    step_id: str
    session_id: str
    message_id: str
    step_type: str  # "classification", "explanation", "permission_check", "follow_up", "referral", "feedback"
    input_text: str
    model_output: str
    approval_status: Optional[str] = None  # "approved", "disapproved", None
    tagging_data: Optional[TaggingRecord] = None
    timestamp: datetime = field(default_factory=datetime.now)
    
    def __post_init__(self):
        """Validate step type."""
        if self.step_type not in INTERACTION_STEP_TYPES:
            raise ValueError(f"Invalid step type: {self.step_type}")
        if self.approval_status and self.approval_status not in ("approved", "disapproved"):
            raise ValueError(f"Invalid approval status: {self.approval_status}")
    
    def to_dict(self) -> dict:
        """Convert log entry to dictionary for serialization."""
        return {
            "step_id": self.step_id,
            "session_id": self.session_id,
            "message_id": self.message_id,
            "step_type": self.step_type,
            "input_text": self.input_text,
            "model_output": self.model_output,
            "approval_status": self.approval_status,
            "tagging_data": self.tagging_data.to_dict() if self.tagging_data else None,
            "timestamp": self.timestamp.isoformat(),
        }
    
    @classmethod
    def from_dict(cls, data: dict) -> "InteractionStepLog":
        """Create log entry from dictionary."""
        data_copy = data.copy()
        if isinstance(data_copy.get("timestamp"), str):
            data_copy["timestamp"] = datetime.fromisoformat(data_copy["timestamp"])
        
        # Convert nested tagging data
        tagging_data = data_copy.pop("tagging_data", None)
        if tagging_data:
            tagging_data = TaggingRecord.from_dict(tagging_data)
        
        log = cls(**data_copy)
        log.tagging_data = tagging_data
        return log