tools/explain_in_layman_language.py

"""
explain_in_layman_language.py
-----------------------------

Tool for explaining medical assessment and plans in patient-friendly language with educational resources.

This tool takes complex medical terminology from assessments and plans and translates them into
easy-to-understand language for patients. It also searches for reliable educational resources
to help patients better understand their conditions and treatments.
"""

import asyncio
from typing import Dict, List, Union, Any
from tools.base import Tool
from tools.utils import ToolExecutionError, logger
from tools.internet_search import InternetSearchTool


class ExplainInLaymanLanguageTool(Tool):
    """
    Tool for explaining medical assessments and plans in patient-friendly language.
    
    This tool:
    1. Takes medical assessment and plan text
    2. Translates complex medical terminology into layman terms
    3. Searches for 2-3 reliable educational links on the topics
    4. Provides a comprehensive patient-friendly explanation
    """
    
    def __init__(self) -> None:
        """Initialize the ExplainInLaymanLanguageTool."""
        super().__init__()
        self.name = "explain_in_layman_language"
        self.description = "Explain medical assessment and plan in patient-friendly language with educational resources"
        self.args_schema = {
            "type": "object",
            "properties": {
                "assessment_and_plan": {
                    "type": "string",
                    "description": "The medical assessment and plan text to explain in layman terms"
                },
                "patient_context": {
                    "type": "string",
                    "description": "Additional context about the patient (age, relevant history, etc.)",
                    "default": ""
                }
            },
            "required": ["assessment_and_plan"]
        }
        
        # Initialize internet search tool for finding educational resources
        self.internet_search = InternetSearchTool()
    
    def openai_spec(self, legacy=False):
        """Return OpenAI function specification for this tool."""
        return {
            "name": self.name,
            "description": self.description,
            "parameters": self.args_schema
        }
    
    async def run(
        self,
        assessment_and_plan: str,
        patient_context: str = ""
    ) -> Dict[str, Any]:
        """
        Explain medical assessment and plan in layman language with educational resources.
        
        Args:
            assessment_and_plan (str): The medical assessment and plan to explain
            patient_context (str): Additional patient context
            
        Returns:
            Dict[str, Any]: Patient-friendly explanation with educational resources
        """
        try:
            logger.info(f"Explaining medical assessment in layman language")
            
            # Step 1: Extract key medical topics from the assessment and plan
            key_topics = self._extract_medical_topics(assessment_and_plan)
            
            # Step 2: Create patient-friendly explanation
            layman_explanation = self._create_layman_explanation(
                assessment_and_plan, 
                patient_context, 
                key_topics
            )
            
            # Step 3: Search for educational resources
            educational_links = await self._find_educational_resources(key_topics)
            
            # Step 4: Combine everything into a comprehensive response
            result = {
                "layman_explanation": layman_explanation,
                "educational_resources": educational_links,
                "key_topics_covered": key_topics,
                "patient_context": patient_context
            }
            
            logger.info(f"Successfully created layman explanation with {len(educational_links)} educational resources")
            return result
            
        except Exception as e:
            logger.error(f"ExplainInLaymanLanguageTool failed: {e}", exc_info=True)
            raise ToolExecutionError(f"Failed to explain assessment in layman language: {e}")
    
    def _extract_medical_topics(self, assessment_and_plan: str) -> List[str]:
        """Extract key medical topics from the assessment and plan."""
        # Common medical terms and their categories
        medical_terms = {
            # Infections
            "pneumonia": "lung infection",
            "sepsis": "blood infection",
            "cellulitis": "skin infection",
            "uti": "urinary tract infection",
            "urinary tract infection": "urinary tract infection",
            "meningitis": "brain infection",
            "endocarditis": "heart infection",
            
            # Antibiotics
            "antibiotic": "antibiotic treatment",
            "vancomycin": "antibiotic treatment",
            "ceftriaxone": "antibiotic treatment",
            "azithromycin": "antibiotic treatment",
            "levofloxacin": "antibiotic treatment",
            "piperacillin": "antibiotic treatment",
            "meropenem": "antibiotic treatment",
            
            # Conditions
            "diabetes": "diabetes",
            "hypertension": "high blood pressure",
            "copd": "chronic lung disease",
            "asthma": "asthma",
            "heart failure": "heart failure",
            "kidney disease": "kidney disease",
            "liver disease": "liver disease",
            
            # Procedures
            "blood culture": "blood tests",
            "chest x-ray": "chest imaging",
            "ct scan": "CT scan",
            "mri": "MRI scan",
            "lumbar puncture": "spinal fluid test",
            "biopsy": "tissue sample",
            
            # Symptoms
            "fever": "fever",
            "shortness of breath": "breathing difficulty",
            "chest pain": "chest pain",
            "abdominal pain": "stomach pain",
            "nausea": "nausea",
            "vomiting": "vomiting",
            "diarrhea": "diarrhea"
        }
        
        topics = []
        assessment_lower = assessment_and_plan.lower()
        
        for term, category in medical_terms.items():
            if term in assessment_lower:
                if category not in topics:
                    topics.append(category)
        
        # If no specific topics found, extract general categories
        if not topics:
            if "infection" in assessment_lower:
                topics.append("infection")
            if "antibiotic" in assessment_lower:
                topics.append("antibiotic treatment")
            if "treatment" in assessment_lower:
                topics.append("treatment")
        
        return topics[:5]  # Limit to top 5 topics
    
    def _create_layman_explanation(
        self, 
        assessment_and_plan: str, 
        patient_context: str, 
        key_topics: List[str]
    ) -> str:
        """Create a patient-friendly explanation of the assessment and plan."""
        
        # Medical terminology translations
        translations = {
            # Infections
            "pneumonia": "a lung infection",
            "sepsis": "a serious blood infection",
            "bacteremia": "bacteria in the blood",
            "cellulitis": "a skin and soft tissue infection",
            "uti": "a urinary tract infection (bladder infection)",
            "urinary tract infection": "a bladder infection",
            "pyelonephritis": "a kidney infection",
            "meningitis": "an infection of the brain and spinal cord lining",
            "endocarditis": "an infection of the heart valves",
            
            # Antibiotics
            "vancomycin": "a strong antibiotic medication",
            "ceftriaxone": "an antibiotic medication",
            "azithromycin": "an antibiotic medication (Z-pack)",
            "levofloxacin": "an antibiotic medication",
            "piperacillin-tazobactam": "a combination antibiotic medication",
            "meropenem": "a powerful antibiotic medication",
            "clindamycin": "an antibiotic medication",
            
            # Medical procedures
            "blood culture": "blood tests to check for bacteria",
            "chest x-ray": "a picture of your lungs",
            "ct scan": "a detailed scan using X-rays",
            "mri": "a detailed scan using magnets",
            "lumbar puncture": "a procedure to test spinal fluid",
            "biopsy": "taking a small sample of tissue for testing",
            
            # Conditions
            "copd": "chronic obstructive pulmonary disease (long-term lung condition)",
            "chf": "congestive heart failure (heart not pumping well)",
            "dm": "diabetes mellitus (diabetes)",
            "htn": "hypertension (high blood pressure)",
            "ckd": "chronic kidney disease (long-term kidney problems)",
            
            # Symptoms
            "dyspnea": "shortness of breath",
            "tachycardia": "fast heart rate",
            "hypotension": "low blood pressure",
            "hypertension": "high blood pressure",
            "pyrexia": "fever",
            "malaise": "feeling unwell",
            "lethargy": "extreme tiredness"
        }
        
        # Start with a friendly introduction
        explanation = "Here's what your medical team has found and what the plan is, explained in simple terms:\n\n"
        
        # Add patient context if provided
        if patient_context:
            explanation += f"**About You:** {patient_context}\n\n"
        
        # Process the assessment and plan
        explanation += "**What's Happening:**\n"
        
        # Simple word replacement for common medical terms
        simplified_text = assessment_and_plan.lower()
        for medical_term, layman_term in translations.items():
            simplified_text = simplified_text.replace(medical_term, layman_term)
        
        # Add the simplified explanation
        explanation += simplified_text.capitalize() + "\n\n"
        
        # Add key points section
        if key_topics:
            explanation += "**Key Points to Remember:**\n"
            for topic in key_topics:
                explanation += f"• {topic.capitalize()}\n"
            explanation += "\n"
        
        # Add encouragement and next steps
        explanation += "**What This Means for You:**\n"
        explanation += "Your medical team is working to give you the best care possible. "
        explanation += "It's important to follow the treatment plan and ask questions if anything is unclear. "
        explanation += "Your health and understanding are our top priorities.\n\n"
        
        return explanation
    
    async def _find_educational_resources(self, key_topics: List[str]) -> List[Dict[str, str]]:
        """Find 2-3 reliable educational resources about the key topics."""
        educational_links = []
        
        # Trusted medical education websites
        trusted_sites = [
            "mayoclinic.org",
            "webmd.com",
            "healthline.com",
            "medlineplus.gov",
            "cdc.gov",
            "nih.gov",
            "patient.info"
        ]
        
        try:
            # Search for educational resources for each topic
            for topic in key_topics[:3]:  # Limit to 3 topics
                search_queries = [
                    f"{topic} patient education site:mayoclinic.org",
                    f"{topic} patient information site:medlineplus.gov",
                    f"what is {topic} patient guide site:healthline.com"
                ]
                
                for query in search_queries:
                    try:
                        search_results = await self.internet_search.run(q=query, max_results=3)
                        
                        # Parse the search results string to extract links
                        if search_results and isinstance(search_results, str):
                            # Extract links from the formatted string response
                            import re
                            # Pattern to match [Read more](url) links
                            link_pattern = r'\[Read more\]\(([^)]+)\)'
                            title_pattern = r'\*\*([^*]+)\*\*'
                            
                            links = re.findall(link_pattern, search_results)
                            titles = re.findall(title_pattern, search_results)
                            
                            # Match titles with links
                            for i, (title, link) in enumerate(zip(titles, links)):
                                if any(site in link.lower() for site in trusted_sites):
                                    educational_links.append({
                                        "title": title,
                                        "url": link,
                                        "topic": topic,
                                        "description": f"Learn more about {topic}"
                                    })
                                    break  # Found a good link for this topic
                        
                        if len(educational_links) >= 3:
                            break
                            
                    except Exception as e:
                        logger.warning(f"Failed to search for {topic}: {e}")
                        continue
                
                if len(educational_links) >= 3:
                    break
            
            # If we don't have enough links, add some general reliable resources
            if len(educational_links) < 2:
                general_resources = [
                    {
                        "title": "MedlinePlus Health Information",
                        "url": "https://medlineplus.gov/",
                        "topic": "general health",
                        "description": "Reliable health information from the National Library of Medicine"
                    },
                    {
                        "title": "Mayo Clinic Patient Education",
                        "url": "https://www.mayoclinic.org/patient-visitor-guide/patient-education",
                        "topic": "general health",
                        "description": "Comprehensive patient education resources from Mayo Clinic"
                    },
                    {
                        "title": "CDC Health Information",
                        "url": "https://www.cdc.gov/healthypeople/",
                        "topic": "general health",
                        "description": "Health information and resources from the CDC"
                    }
                ]
                
                # Add general resources to reach at least 2-3 links
                for resource in general_resources:
                    if len(educational_links) < 3:
                        educational_links.append(resource)
            
        except Exception as e:
            logger.warning(f"Failed to find educational resources: {e}")
            # Provide fallback resources
            educational_links = [
                {
                    "title": "MedlinePlus - Easy-to-Read Health Information",
                    "url": "https://medlineplus.gov/",
                    "topic": "general health",
                    "description": "Trusted health information from the National Library of Medicine"
                },
                {
                    "title": "Mayo Clinic - Patient Education",
                    "url": "https://www.mayoclinic.org/diseases-conditions",
                    "topic": "general health",
                    "description": "Comprehensive information about diseases and conditions"
                }
            ]
        
        return educational_links[:3]  # Return maximum 3 links