"""
Medical Prompt Templates for MedGemma Synthesis
Comprehensive templates for generating clinician-level and patient-friendly summaries

Author: MiniMax Agent
Date: 2025-10-29
Version: 1.0.0
"""

from typing import Dict, Any, List, Optional
from enum import Enum


class SummaryType(Enum):
    """Types of medical summaries that can be generated"""
    CLINICIAN_TECHNICAL = "clinician_technical"
    PATIENT_FRIENDLY = "patient_friendly"
    MULTI_MODAL = "multi_modal"
    RISK_ASSESSMENT = "risk_assessment"


class PromptTemplateLibrary:
    """
    Comprehensive library of medical prompt templates for MedGemma
    Supports all medical modalities with evidence-based generation
    """
    
    @staticmethod
    def get_clinician_summary_template(
        modality: str,
        structured_data: Dict[str, Any],
        model_outputs: List[Dict[str, Any]],
        confidence_scores: Dict[str, float]
    ) -> str:
        """
        Generate clinician-level technical summary prompt
        
        Features:
        - Technical medical terminology
        - Detailed analysis with evidence
        - Confidence scores and uncertainty
        - Clinical decision support
        """
        
        if modality == "ECG":
            return PromptTemplateLibrary._ecg_clinician_template(
                structured_data, model_outputs, confidence_scores
            )
        elif modality == "radiology":
            return PromptTemplateLibrary._radiology_clinician_template(
                structured_data, model_outputs, confidence_scores
            )
        elif modality == "laboratory":
            return PromptTemplateLibrary._laboratory_clinician_template(
                structured_data, model_outputs, confidence_scores
            )
        elif modality == "clinical_notes":
            return PromptTemplateLibrary._clinical_notes_clinician_template(
                structured_data, model_outputs, confidence_scores
            )
        else:
            return PromptTemplateLibrary._general_clinician_template(
                structured_data, model_outputs, confidence_scores
            )
    
    @staticmethod
    def get_patient_summary_template(
        modality: str,
        structured_data: Dict[str, Any],
        model_outputs: List[Dict[str, Any]],
        confidence_scores: Dict[str, float]
    ) -> str:
        """
        Generate patient-friendly summary prompt
        
        Features:
        - Plain language explanations
        - Key findings highlighted
        - Actionable next steps
        - Reassurance when appropriate
        """
        
        if modality == "ECG":
            return PromptTemplateLibrary._ecg_patient_template(
                structured_data, model_outputs, confidence_scores
            )
        elif modality == "radiology":
            return PromptTemplateLibrary._radiology_patient_template(
                structured_data, model_outputs, confidence_scores
            )
        elif modality == "laboratory":
            return PromptTemplateLibrary._laboratory_patient_template(
                structured_data, model_outputs, confidence_scores
            )
        elif modality == "clinical_notes":
            return PromptTemplateLibrary._clinical_notes_patient_template(
                structured_data, model_outputs, confidence_scores
            )
        else:
            return PromptTemplateLibrary._general_patient_template(
                structured_data, model_outputs, confidence_scores
            )
    
    # ========================
    # ECG TEMPLATES
    # ========================
    
    @staticmethod
    def _ecg_clinician_template(
        data: Dict[str, Any],
        outputs: List[Dict[str, Any]],
        confidence: Dict[str, float]
    ) -> str:
        """Clinician-level ECG summary template"""
        
        intervals = data.get("intervals", {})
        rhythm = data.get("rhythm_classification", {})
        arrhythmia_probs = data.get("arrhythmia_probabilities", {})
        derived = data.get("derived_features", {})
        
        overall_confidence = confidence.get("overall_confidence", 0.0)
        
        prompt = f"""You are a medical AI assistant generating a comprehensive ECG analysis report for clinicians.

PATIENT CONTEXT:
- Document ID: {data.get('metadata', {}).get('document_id', 'N/A')}
- Facility: {data.get('metadata', {}).get('facility', 'N/A')}
- Recording Date: {data.get('metadata', {}).get('document_date', 'N/A')}

ECG MEASUREMENTS:
- Heart Rate: {rhythm.get('heart_rate_bpm', 'N/A')} bpm
- PR Interval: {intervals.get('pr_ms', 'N/A')} ms
- QRS Duration: {intervals.get('qrs_ms', 'N/A')} ms
- QT Interval: {intervals.get('qt_ms', 'N/A')} ms
- QTc Interval: {intervals.get('qtc_ms', 'N/A')} ms
- RR Interval: {intervals.get('rr_ms', 'N/A')} ms

RHYTHM ANALYSIS:
- Primary Rhythm: {rhythm.get('primary_rhythm', 'N/A')}
- Rhythm Regularity: {rhythm.get('heart_rate_regularity', 'N/A')}
- Detected Arrhythmias: {', '.join(rhythm.get('arrhythmia_types', [])) or 'None'}

ARRHYTHMIA PROBABILITIES:
- Normal Sinus Rhythm: {arrhythmia_probs.get('normal_rhythm', 'N/A')}
- Atrial Fibrillation: {arrhythmia_probs.get('atrial_fibrillation', 'N/A')}
- Atrial Flutter: {arrhythmia_probs.get('atrial_flutter', 'N/A')}
- Ventricular Tachycardia: {arrhythmia_probs.get('ventricular_tachycardia', 'N/A')}
- Heart Block: {arrhythmia_probs.get('heart_block', 'N/A')}

ST-SEGMENT & T-WAVE FINDINGS:
- ST Elevation: {derived.get('st_elevation_mm', 'None detected')}
- ST Depression: {derived.get('st_depression_mm', 'None detected')}
- T-wave Abnormalities: {', '.join(derived.get('t_wave_abnormalities', [])) or 'None'}
- Axis Deviation: {derived.get('axis_deviation', 'Normal')}

AI MODEL OUTPUTS:
{PromptTemplateLibrary._format_model_outputs(outputs)}

ANALYSIS CONFIDENCE: {overall_confidence * 100:.1f}%

INSTRUCTIONS:
Generate a comprehensive clinical ECG report with the following sections:

1. TECHNICAL SUMMARY
   - Concise interpretation of rhythm and intervals
   - Significance of any abnormal findings
   
2. CLINICAL SIGNIFICANCE
   - Pathophysiological implications
   - Risk stratification (low/moderate/high)
   
3. DIFFERENTIAL DIAGNOSIS
   - Most likely diagnoses based on findings
   - Alternative considerations
   
4. RECOMMENDATIONS
   - Immediate actions required (if any)
   - Follow-up studies or monitoring
   - Cardiology referral if indicated
   
5. CONFIDENCE EXPLANATION
   - Why the AI confidence is {overall_confidence * 100:.1f}%
   - Which findings are most/least certain
   - Limitations of the analysis

Use precise medical terminology. Be evidence-based. Flag any critical findings requiring immediate attention.

Generate the report now:"""
        
        return prompt
    
    @staticmethod
    def _ecg_patient_template(
        data: Dict[str, Any],
        outputs: List[Dict[str, Any]],
        confidence: Dict[str, float]
    ) -> str:
        """Patient-friendly ECG summary template"""
        
        rhythm = data.get("rhythm_classification", {})
        intervals = data.get("intervals", {})
        
        prompt = f"""You are a medical AI assistant explaining ECG results to a patient in simple, clear language.

YOUR ECG RESULTS:
- Heart Rate: {rhythm.get('heart_rate_bpm', 'N/A')} beats per minute
- Heart Rhythm: {rhythm.get('primary_rhythm', 'N/A')}

WHAT THIS MEANS:
Generate a patient-friendly explanation that:

1. WHAT WE FOUND
   - Explain the heart rate and rhythm in simple terms
   - Describe any abnormalities without medical jargon
   
2. WHAT THIS MEANS FOR YOU
   - Is this normal or concerning?
   - What might be causing any abnormalities?
   
3. NEXT STEPS
   - What should you do next?
   - Do you need to see a doctor urgently?
   - Any lifestyle changes to consider?
   
4. OUR CONFIDENCE
   - How certain are we about these findings?
   - Why you should still talk to your doctor

Use everyday language. Be reassuring when appropriate. Be clear about urgency if there are concerns.

Generate the patient explanation now:"""
        
        return prompt
    
    # ========================
    # RADIOLOGY TEMPLATES
    # ========================
    
    @staticmethod
    def _radiology_clinician_template(
        data: Dict[str, Any],
        outputs: List[Dict[str, Any]],
        confidence: Dict[str, float]
    ) -> str:
        """Clinician-level radiology summary template"""
        
        findings = data.get("findings", {})
        metrics = data.get("metrics", {})
        images = data.get("image_references", [])
        
        prompt = f"""You are a radiologist AI assistant generating a comprehensive imaging report.

IMAGING STUDY DETAILS:
- Modality: {', '.join([img.get('modality', 'N/A') for img in images[:3]])}
- Body Parts: {', '.join([img.get('body_part', 'N/A') for img in images[:3]])}
- Study Date: {data.get('metadata', {}).get('document_date', 'N/A')}

FINDINGS:
{findings.get('findings_text', 'N/A')}

IMPRESSION:
{findings.get('impression_text', 'N/A')}

CRITICAL FINDINGS: {', '.join(findings.get('critical_findings', [])) or 'None'}
INCIDENTAL FINDINGS: {', '.join(findings.get('incidental_findings', [])) or 'None'}

QUANTITATIVE METRICS:
- Organ Volumes: {metrics.get('organ_volumes', {})}
- Lesion Measurements: {len(metrics.get('lesion_measurements', []))} lesions measured

AI MODEL ANALYSIS:
{PromptTemplateLibrary._format_model_outputs(outputs)}

ANALYSIS CONFIDENCE: {confidence.get('overall_confidence', 0.0) * 100:.1f}%

Generate a structured radiology report with:

1. TECHNIQUE & COMPARISON
2. FINDINGS (organized by anatomical region)
3. IMPRESSION
4. RECOMMENDATIONS
5. CONFIDENCE ASSESSMENT

Use standard radiology terminology (BI-RADS, Lung-RADS, etc. if applicable).

Generate the report now:"""
        
        return prompt
    
    @staticmethod
    def _radiology_patient_template(
        data: Dict[str, Any],
        outputs: List[Dict[str, Any]],
        confidence: Dict[str, float]
    ) -> str:
        """Patient-friendly radiology summary template"""
        
        findings = data.get("findings", {})
        images = data.get("image_references", [])
        
        prompt = f"""You are explaining imaging results to a patient in clear, simple language.

YOUR IMAGING STUDY:
- Type of Scan: {', '.join([img.get('modality', 'N/A') for img in images[:3]])}
- Body Area: {', '.join([img.get('body_part', 'N/A') for img in images[:3]])}

Generate a patient-friendly explanation:

1. WHAT THE SCAN SHOWED
   - Main findings in simple terms
   - Any areas of concern
   
2. WHAT THIS MEANS
   - Are the findings normal or abnormal?
   - What conditions might this suggest?
   
3. NEXT STEPS
   - Do you need additional tests?
   - Should you see a specialist?
   - Timeline for follow-up
   
4. QUESTIONS TO ASK YOUR DOCTOR
   - List 3-4 relevant questions

Use everyday language. Explain medical terms when necessary. Be clear about urgency.

Generate the patient explanation now:"""
        
        return prompt
    
    # ========================
    # LABORATORY TEMPLATES
    # ========================
    
    @staticmethod
    def _laboratory_clinician_template(
        data: Dict[str, Any],
        outputs: List[Dict[str, Any]],
        confidence: Dict[str, float]
    ) -> str:
        """Clinician-level laboratory results template"""
        
        tests = data.get("tests", [])
        abnormal_count = data.get("abnormal_count", 0)
        critical_values = data.get("critical_values", [])
        
        test_summary = "\n".join([
            f"- {test.get('test_name', 'N/A')}: {test.get('value', 'N/A')} {test.get('unit', '')} "
            f"(Ref: {test.get('reference_range_low', 'N/A')}-{test.get('reference_range_high', 'N/A')}) "
            f"{test.get('flags', [])}"
            for test in tests[:20]  # Limit to 20 tests
        ])
        
        prompt = f"""You are a clinical laboratory AI assistant generating a comprehensive lab results analysis.

LABORATORY PANEL:
- Panel Type: {data.get('panel_name', 'General Laboratory Panel')}
- Collection Date: {data.get('collection_date', 'N/A')}
- Total Tests: {len(tests)}
- Abnormal Results: {abnormal_count}
- Critical Values: {len(critical_values)}

TEST RESULTS:
{test_summary}

CRITICAL VALUES: {', '.join(critical_values) or 'None'}

AI MODEL ANALYSIS:
{PromptTemplateLibrary._format_model_outputs(outputs)}

ANALYSIS CONFIDENCE: {confidence.get('overall_confidence', 0.0) * 100:.1f}%

Generate a comprehensive laboratory interpretation with:

1. SUMMARY OF KEY FINDINGS
   - Normal vs abnormal results
   - Critical values requiring immediate attention
   
2. CLINICAL CORRELATION
   - Pattern recognition (e.g., renal dysfunction, electrolyte imbalance)
   - Physiological significance
   
3. DIFFERENTIAL DIAGNOSIS
   - Most likely conditions based on lab pattern
   
4. RECOMMENDATIONS
   - Immediate interventions for critical values
   - Additional testing needed
   - Follow-up timeline
   
5. CONFIDENCE ASSESSMENT
   - Reliability of each test result
   - Need for repeat testing

Generate the interpretation now:"""
        
        return prompt
    
    @staticmethod
    def _laboratory_patient_template(
        data: Dict[str, Any],
        outputs: List[Dict[str, Any]],
        confidence: Dict[str, float]
    ) -> str:
        """Patient-friendly laboratory results template"""
        
        tests = data.get("tests", [])
        abnormal_count = data.get("abnormal_count", 0)
        
        prompt = f"""You are explaining laboratory test results to a patient in simple language.

YOUR LAB RESULTS:
- Total Tests: {len(tests)}
- Abnormal Results: {abnormal_count}

Generate a patient-friendly explanation:

1. OVERVIEW
   - What tests were done and why
   - Overall picture (mostly normal, some concerns, etc.)
   
2. KEY FINDINGS
   - Which results are normal
   - Which results are outside the normal range
   - What each abnormal result means in simple terms
   
3. WHAT THIS MEANS FOR YOUR HEALTH
   - Are these results concerning?
   - What conditions might they suggest?
   
4. NEXT STEPS
   - Do you need to see your doctor urgently?
   - Lifestyle changes that might help
   - Additional tests that might be needed
   
5. IMPORTANT NOTES
   - Lab values can vary based on many factors
   - Always discuss results with your doctor

Use everyday language. Explain abbreviations. Be clear about urgency.

Generate the patient explanation now:"""
        
        return prompt
    
    # ========================
    # CLINICAL NOTES TEMPLATES
    # ========================
    
    @staticmethod
    def _clinical_notes_clinician_template(
        data: Dict[str, Any],
        outputs: List[Dict[str, Any]],
        confidence: Dict[str, float]
    ) -> str:
        """Clinician-level clinical notes summary template"""
        
        sections = data.get("sections", [])
        entities = data.get("entities", [])
        diagnoses = data.get("diagnoses", [])
        medications = data.get("medications", [])
        
        sections_summary = "\n".join([
            f"- {section.get('section_type', 'N/A')}: {section.get('content', 'N/A')[:200]}..."
            for section in sections[:10]
        ])
        
        prompt = f"""You are a clinical documentation AI assistant synthesizing medical notes.

NOTE TYPE: {data.get('note_type', 'Clinical Documentation')}
DOCUMENTATION DATE: {data.get('metadata', {}).get('document_date', 'N/A')}

CLINICAL SECTIONS:
{sections_summary}

EXTRACTED ENTITIES:
- Diagnoses: {', '.join(diagnoses[:10]) or 'None identified'}
- Medications: {', '.join(medications[:10]) or 'None identified'}

AI MODEL ANALYSIS:
{PromptTemplateLibrary._format_model_outputs(outputs)}

ANALYSIS CONFIDENCE: {confidence.get('overall_confidence', 0.0) * 100:.1f}%

Generate a comprehensive clinical synthesis with:

1. CLINICAL SUMMARY
   - Chief complaint and HPI synthesis
   - Pertinent positives and negatives
   
2. ASSESSMENT
   - Problem list with prioritization
   - Clinical reasoning
   
3. PLAN
   - Management for each problem
   - Medications and interventions
   - Follow-up and monitoring
   
4. DOCUMENTATION QUALITY
   - Completeness assessment
   - Missing information
   
5. CONFIDENCE ASSESSMENT

Generate the clinical synthesis now:"""
        
        return prompt
    
    @staticmethod
    def _clinical_notes_patient_template(
        data: Dict[str, Any],
        outputs: List[Dict[str, Any]],
        confidence: Dict[str, float]
    ) -> str:
        """Patient-friendly clinical notes summary template"""
        
        diagnoses = data.get("diagnoses", [])
        medications = data.get("medications", [])
        
        prompt = f"""You are explaining a clinical visit summary to a patient in clear, simple language.

Generate a patient-friendly visit summary:

1. REASON FOR YOUR VISIT
   - Why you came to see the doctor
   
2. WHAT THE DOCTOR FOUND
   - Key findings from examination
   - Test results discussed
   
3. YOUR DIAGNOSES
   - {', '.join(diagnoses[:5]) if diagnoses else 'To be discussed with your doctor'}
   - What each diagnosis means in simple terms
   
4. YOUR TREATMENT PLAN
   - Medications prescribed
   - Other treatments or therapies
   
5. WHAT YOU NEED TO DO
   - Follow-up appointments
   - Tests or procedures needed
   - Lifestyle changes
   - Warning signs to watch for
   
6. QUESTIONS FOR YOUR DOCTOR
   - List important questions to ask

Use everyday language. Explain medical terms. Organize by priority.

Generate the patient summary now:"""
        
        return prompt
    
    # ========================
    # GENERAL TEMPLATES
    # ========================
    
    @staticmethod
    def _general_clinician_template(
        data: Dict[str, Any],
        outputs: List[Dict[str, Any]],
        confidence: Dict[str, float]
    ) -> str:
        """General clinician-level summary template"""
        
        prompt = f"""You are a medical AI assistant generating a comprehensive clinical summary.

DOCUMENT TYPE: {data.get('metadata', {}).get('source_type', 'Medical Document')}
DOCUMENT DATE: {data.get('metadata', {}).get('document_date', 'N/A')}

AI MODEL ANALYSIS:
{PromptTemplateLibrary._format_model_outputs(outputs)}

ANALYSIS CONFIDENCE: {confidence.get('overall_confidence', 0.0) * 100:.1f}%

Generate a structured medical summary with:
1. KEY FINDINGS
2. CLINICAL SIGNIFICANCE
3. RECOMMENDATIONS
4. CONFIDENCE ASSESSMENT

Use appropriate medical terminology.

Generate the summary now:"""
        
        return prompt
    
    @staticmethod
    def _general_patient_template(
        data: Dict[str, Any],
        outputs: List[Dict[str, Any]],
        confidence: Dict[str, float]
    ) -> str:
        """General patient-friendly summary template"""
        
        prompt = f"""You are explaining medical information to a patient in simple, clear language.

Generate a patient-friendly explanation:
1. WHAT WE FOUND
2. WHAT THIS MEANS FOR YOU
3. NEXT STEPS
4. QUESTIONS TO ASK YOUR DOCTOR

Use everyday language. Be clear and reassuring when appropriate.

Generate the explanation now:"""
        
        return prompt
    
    # ========================
    # MULTI-MODAL SYNTHESIS
    # ========================
    
    @staticmethod
    def get_multi_modal_synthesis_template(
        modalities: List[str],
        all_data: Dict[str, Dict[str, Any]],
        confidence_scores: Dict[str, float]
    ) -> str:
        """
        Generate prompt for multi-modal clinical synthesis
        Combines multiple document types into unified summary
        """
        
        modality_summaries = []
        for modality in modalities:
            data = all_data.get(modality, {})
            modality_summaries.append(f"- {modality.upper()}: Available with {confidence_scores.get(modality, 0.0)*100:.1f}% confidence")
        
        prompt = f"""You are a medical AI assistant synthesizing multiple medical documents into a comprehensive clinical picture.

AVAILABLE DOCUMENTS:
{chr(10).join(modality_summaries)}

TASK:
Generate a unified clinical summary that:

1. INTEGRATED CLINICAL PICTURE
   - Synthesize findings across all modalities
   - Identify consistent patterns
   - Flag contradictions or discrepancies
   
2. TIMELINE CORRELATION
   - How findings relate temporally
   - Disease progression or improvement
   
3. COMPREHENSIVE ASSESSMENT
   - Overall patient status
   - Risk stratification
   
4. COORDINATED CARE PLAN
   - Unified recommendations
   - Priority actions
   - Specialist referrals
   
5. CONFIDENCE SYNTHESIS
   - Overall reliability of the integrated analysis
   - Areas needing additional investigation

Generate the integrated clinical synthesis now:"""
        
        return prompt
    
    # ========================
    # UTILITY METHODS
    # ========================
    
    @staticmethod
    def _format_model_outputs(outputs: List[Dict[str, Any]]) -> str:
        """Format model outputs for inclusion in prompts"""
        if not outputs:
            return "No specialized model outputs available"
        
        formatted = []
        for idx, output in enumerate(outputs[:5], 1):  # Limit to top 5
            model_name = output.get("model_name", "Unknown Model")
            domain = output.get("domain", "general")
            result = output.get("result", {})
            
            # Extract key information from result
            if isinstance(result, dict):
                confidence = result.get("confidence", 0.0)
                summary = result.get("summary", result.get("analysis", "Analysis completed"))[:200]
                formatted.append(f"{idx}. {model_name} ({domain}): {summary}... [Confidence: {confidence*100:.1f}%]")
            else:
                formatted.append(f"{idx}. {model_name} ({domain}): {str(result)[:200]}...")
        
        return "\n".join(formatted)
    
    @staticmethod
    def get_confidence_explanation_template(
        confidence_scores: Dict[str, float],
        modality: str
    ) -> str:
        """Generate prompt for explaining confidence scores"""
        
        overall = confidence_scores.get("overall_confidence", 0.0)
        extraction = confidence_scores.get("extraction_confidence", 0.0)
        model = confidence_scores.get("model_confidence", 0.0)
        quality = confidence_scores.get("data_quality", 0.0)
        
        if overall >= 0.85:
            threshold = "AUTO-APPROVED (≥85%)"
        elif overall >= 0.60:
            threshold = "REQUIRES REVIEW (60-85%)"
        else:
            threshold = "MANUAL REVIEW REQUIRED (<60%)"
        
        prompt = f"""Explain the confidence scores for this {modality} analysis to a clinician:

CONFIDENCE BREAKDOWN:
- Overall Confidence: {overall*100:.1f}% [{threshold}]
- Data Extraction: {extraction*100:.1f}%
- Model Analysis: {model*100:.1f}%
- Data Quality: {quality*100:.1f}%

Generate a brief explanation that:
1. Why this confidence level?
2. What factors contributed to the score?
3. What should the clinician be aware of?
4. Is human review recommended?

Be concise and practical.

Generate the explanation now:"""
        
        return prompt