response_formatter.py

# response_formatter.py
from typing import Dict, List, Any, Tuple
import re

class ResponseFormatter:
    @staticmethod
    def format_healthcare_response(scenario_text: str, analysis_results: Dict[str, Any]) -> str:
        """Format healthcare analysis response based on scenario tasks"""
        # Extract tasks from scenario
        tasks = ResponseFormatter._extract_scenario_tasks(scenario_text)
        
        response = "# Healthcare Scenario Analysis\n\n"
        
        # Executive Summary
        response += "## Executive Summary\n\n"
        response += ResponseFormatter._generate_executive_summary(analysis_results)
        response += "\n\n"
        
        # Process each task
        for task in tasks:
            section_content = ResponseFormatter._process_task(task, analysis_results)
            if section_content:
                response += section_content + "\n\n"
        
        # Validation Results
        if "validation" in analysis_results:
            response += "## Analysis Validation\n\n"
            response += ResponseFormatter._format_validation_results(analysis_results["validation"])
            response += "\n\n"
        
        # Provenance
        response += "## Provenance\n\n"
        response += "This analysis is based on:\n"
        response += "- Scenario description provided by the user\n"
        response += "- Uploaded data files\n"
        response += "- Calculations performed on the provided data\n"
        
        return response
    
    @staticmethod
    def _extract_scenario_tasks(scenario_text: str) -> List[str]:
        """Extract specific tasks from scenario text"""
        tasks = []
        lines = scenario_text.split('\n')
        in_tasks = False
        current_task = ""
        
        for line in lines:
            line = line.strip()
            
            # Check if we're entering the tasks section
            if line.lower().startswith('tasks'):
                in_tasks = True
                continue
            
            # Check if we're leaving the tasks section
            if in_tasks and (line.lower().startswith('operational recommendations') or 
                            line.lower().startswith('future integration') or
                            line.lower().startswith('prepare your findings')):
                in_tasks = False
                if current_task:
                    tasks.append(current_task)
                continue
            
            # Collect task content
            if in_tasks:
                if line and (line.startswith(('1.', '2.', '3.', '4.', '5.', '6.', '7.', '8.', '9.')) or 
                           line.startswith(('•', '-', '*'))):
                    if current_task:
                        tasks.append(current_task)
                    current_task = line
                elif line:
                    current_task += " " + line
        
        # Add the last task if exists
        if current_task:
            tasks.append(current_task)
            
        return tasks
    
    @staticmethod
    def _process_task(task: str, analysis_results: Dict[str, Any]) -> str:
        """Process a single task and return formatted content"""
        task_lower = task.lower()
        
        # Data Preparation Task
        if "data preparation" in task_lower or "load the data" in task_lower:
            return ResponseFormatter._handle_data_preparation_task(task, analysis_results)
        
        # Bed Capacity Analysis Tasks
        elif "bed capacity" in task_lower or "summarize bed capacity" in task_lower:
            return ResponseFormatter._handle_bed_capacity_task(task, analysis_results)
        
        elif "facilities with the greatest declines" in task_lower:
            return ResponseFormatter._handle_facility_declines_task(task, analysis_results)
        
        elif "long-term care capacity" in task_lower or "assess long-term care" in task_lower:
            return ResponseFormatter._handle_long_term_care_task(task, analysis_results)
        
        # Operational Recommendations Tasks
        elif "operational recommendations" in task_lower or "recommend actions" in task_lower:
            return ResponseFormatter._handle_recommendations_task(task, analysis_results)
        
        elif "future integration" in task_lower or "augmented ai" in task_lower:
            return ResponseFormatter._handle_integration_task(task, analysis_results)
        
        # Generic task handler
        else:
            return ResponseFormatter._handle_generic_task(task, analysis_results)
    
    @staticmethod
    def _handle_data_preparation_task(task: str, analysis_results: Dict[str, Any]) -> str:
        """Handle data preparation task"""
        response = "## Data Preparation\n\n"
        
        if "data_preparation" not in analysis_results:
            return response + "No data preparation results available"
        
        data_prep = analysis_results["data_preparation"]
        
        # Facility type frequency
        if "facility_type_frequency" in data_prep:
            response += "### Facility Type Frequency\n\n"
            response += "| Facility Type | Count |\n"
            response += "|---------------|-------|\n"
            for ftype, count in data_prep["facility_type_frequency"].items():
                response += f"| {ftype} | {count} |\n"
            response += "\n"
        
        # Top cities with facility breakdown
        if "top_cities" in data_prep and "city_facility_breakdown" in data_prep:
            response += "### Top Cities by Facility Count\n\n"
            response += "| City | Hospitals | Nursing/Residential | Ambulatory | Total |\n"
            response += "|------|-----------|-------------------|------------|-------|\n"
            
            for city in data_prep["top_cities"]:
                if city in data_prep["city_facility_breakdown"]:
                    breakdown = data_prep["city_facility_breakdown"][city]
                    hospitals = breakdown.get("Hospitals", 0)
                    nursing = breakdown.get("Nursing and residential care facilities", 0)
                    ambulatory = breakdown.get("Ambulatory health care services", 0)
                    total = hospitals + nursing + ambulatory
                    response += f"| {city} | {hospitals} | {nursing} | {ambulatory} | {total} |\n"
            
            response += "\n"
        
        return response
    
    @staticmethod
    def _handle_bed_capacity_task(task: str, analysis_results: Dict[str, Any]) -> str:
        """Handle bed capacity analysis task"""
        response = "## Bed Capacity Analysis\n\n"
        
        if "capacity_analysis" not in analysis_results:
            return response + "No capacity analysis results available"
        
        capacity = analysis_results["capacity_analysis"]
        
        # Bed capacity by zone
        if "bed_capacity_by_zone" in capacity:
            response += "### Bed Capacity by Zone\n\n"
            response += "| Zone | Current Beds | Previous Beds | Absolute Change | Percentage Change |\n"
            response += "|------|--------------|---------------|-----------------|-------------------|\n"
            for zone, data in capacity["bed_capacity_by_zone"].items():
                current = data.get("beds_current", 0)
                previous = data.get("beds_prev", 0)
                abs_change = current - previous
                perc_change = (abs_change / previous * 100) if previous > 0 else 0
                response += f"| {zone} | {current:,} | {previous:,} | {abs_change:+,} | {perc_change:+.1f}% |\n"
            response += "\n"
        
        # Zones with largest decreases
        if "zone_with_largest_absolute_decrease" in capacity:
            response += f"**Zone with Largest Absolute Decrease**: {capacity['zone_with_largest_absolute_decrease']}\n\n"
        
        if "zone_with_largest_percentage_decrease" in capacity:
            response += f"**Zone with Largest Percentage Decrease**: {capacity['zone_with_largest_percentage_decrease']}\n\n"
        
        return response
    
    @staticmethod
    def _handle_facility_declines_task(task: str, analysis_results: Dict[str, Any]) -> str:
        """Handle facilities with greatest declines task"""
        response = "## Facilities with Greatest Declines\n\n"
        
        if "capacity_analysis" not in analysis_results or "facilities_with_declines" not in analysis_results["capacity_analysis"]:
            return response + "No facility decline data available"
        
        facilities = analysis_results["capacity_analysis"]["facilities_with_declines"]
        
        response += "| Facility Name | Zone | Teaching Status | Bed Change |\n"
        response += "|---------------|------|----------------|------------|\n"
        for facility in facilities[:5]:  # Top 5
            name = facility.get("facility_name", "")
            zone = facility.get("zone", "")
            teaching = facility.get("teaching_status", "")
            change = facility.get("bed_change", 0)
            response += f"| {name} | {zone} | {teaching} | {change:+,} |\n"
        
        return response
    
    @staticmethod
    def _handle_long_term_care_task(task: str, analysis_results: Dict[str, Any]) -> str:
        """Handle long-term care capacity assessment task"""
        response = "## Long-Term Care Capacity Assessment\n\n"
        
        if "long_term_care_assessment" not in analysis_results:
            return response + "No long-term care assessment results available"
        
        ltc = analysis_results["long_term_care_assessment"]
        
        # Zone with largest percentage decrease
        if "zone_with_largest_percentage_decrease" in ltc:
            response += f"**Zone with Largest Percentage Decrease**: {ltc['zone_with_largest_percentage_decrease']}\n\n"
        
        # Major city in that zone
        if "major_city_in_zone" in ltc:
            response += f"**Major City in Zone**: {ltc['major_city_in_zone']}\n\n"
        
        # Facility counts in major city
        if "facility_counts" in ltc:
            counts = ltc["facility_counts"]
            response += "### Facility Counts in Major City\n\n"
            response += f"- Hospitals: {counts.get('hospitals', 0)}\n"
            response += f"- Nursing/Residential Care Facilities: {counts.get('nursing_residential_care', 0)}\n"
            response += f"- Ambulatory Health Services: {counts.get('ambulatory', 0)}\n\n"
        
        # Nursing to hospital ratio
        if "nursing_to_hospital_ratio" in ltc:
            ratio = ltc["nursing_to_hospital_ratio"]
            response += f"**Nursing/Residential to Hospital Ratio**: {ratio:.2f}\n\n"
        
        # Capacity assessment
        if "capacity_assessment" in ltc:
            assessment = ltc["capacity_assessment"]
            response += f"**Long-Term Care Capacity Assessment**: {assessment.upper()}\n\n"
            
            if assessment == "insufficient":
                response += "The ratio of nursing/residential care facilities to hospitals is below the recommended threshold of 1.5. "
                response += "This indicates insufficient long-term care capacity to support the healthcare system in this area.\n\n"
            else:
                response += "The ratio of nursing/residential care facilities to hospitals meets or exceeds the recommended threshold. "
                response += "This indicates sufficient long-term care capacity.\n\n"
        
        return response
    
    @staticmethod
    def _handle_recommendations_task(task: str, analysis_results: Dict[str, Any]) -> str:
        """Handle operational recommendations task"""
        response = "## Operational Recommendations\n\n"
        
        if "recommendations" not in analysis_results or not analysis_results["recommendations"]:
            return response + "No recommendations available"
        
        recommendations = analysis_results["recommendations"]
        
        # Group by priority
        high_priority = [r for r in recommendations if r.get("priority") == "High"]
        medium_priority = [r for r in recommendations if r.get("priority") == "Medium"]
        low_priority = [r for r in recommendations if r.get("priority") == "Low"]
        
        if high_priority:
            response += "### High Priority Recommendations\n\n"
            for i, rec in enumerate(high_priority, 1):
                response += f"**{i}. {rec.get('title', 'Recommendation')}**\n\n"
                response += f"{rec.get('description', '')}\n\n"
                if "justification" in rec:
                    response += f"**Justification**: {rec['justification']}\n\n"
                if "expected_impact" in rec:
                    response += f"**Expected Impact**: {rec['expected_impact']}\n\n"
                response += "---\n\n"
        
        if medium_priority:
            response += "### Medium Priority Recommendations\n\n"
            for i, rec in enumerate(medium_priority, 1):
                response += f"**{i}. {rec.get('title', 'Recommendation')}**\n\n"
                response += f"{rec.get('description', '')}\n\n"
                if "justification" in rec:
                    response += f"**Justification**: {rec['justification']}\n\n"
                if "expected_impact" in rec:
                    response += f"**Expected Impact**: {rec['expected_impact']}\n\n"
                response += "---\n\n"
        
        if low_priority:
            response += "### Low Priority Recommendations\n\n"
            for i, rec in enumerate(low_priority, 1):
                response += f"**{i}. {rec.get('title', 'Recommendation')}**\n\n"
                response += f"{rec.get('description', '')}\n\n"
                if "justification" in rec:
                    response += f"**Justification**: {rec['justification']}\n\n"
                if "expected_impact" in rec:
                    response += f"**Expected Impact**: {rec['expected_impact']}\n\n"
                response += "---\n\n"
        
        return response
    
    @staticmethod
    def _handle_integration_task(task: str, analysis_results: Dict[str, Any]) -> str:
        """Handle future integration opportunities task"""
        response = "## Future Integration Opportunities\n\n"
        
        if "future_integration" not in analysis_results:
            return response + "No integration opportunities identified"
        
        integration = analysis_results["future_integration"]
        
        # Summary
        if "summary" in integration:
            response += f"{integration['summary']}\n\n"
        
        # Opportunities
        if "opportunities" in integration:
            for i, opp in enumerate(integration["opportunities"], 1):
                response += f"**{i}. {opp.get('title', 'Opportunity')}**\n\n"
                response += f"{opp.get('description', '')}\n\n"
                
                if "benefits" in opp:
                    response += "**Benefits:**\n"
                    for benefit in opp["benefits"]:
                        response += f"- {benefit}\n"
                    response += "\n"
                
                if "challenges" in opp:
                    response += "**Challenges:**\n"
                    for challenge in opp["challenges"]:
                        response += f"- {challenge}\n"
                    response += "\n"
                
                if "example_metric_or_model" in opp:
                    response += f"**Example Metric/Model**: {opp['example_metric_or_model']}\n\n"
                
                if "timeline" in opp:
                    response += f"**Timeline**: {opp['timeline']}\n\n"
                
                response += "---\n\n"
        
        return response
    
    @staticmethod
    def _handle_generic_task(task: str, analysis_results: Dict[str, Any]) -> str:
        """Handle generic task"""
        # Try to match task to any available analysis results
        task_lower = task.lower()
        
        if "distribution" in task_lower and "facility_distribution" in analysis_results:
            return "## Facility Distribution Analysis\n\n" + ResponseFormatter._format_facility_distribution(analysis_results["facility_distribution"])
        
        if "resource" in task_lower and "resource_allocation" in analysis_results:
            return "## Resource Allocation Analysis\n\n" + ResponseFormatter._format_resource_allocation(analysis_results["resource_allocation"])
        
        if "trend" in task_lower and "trends" in analysis_results:
            return "## Trend Analysis\n\n" + ResponseFormatter._format_trends(analysis_results["trends"])
        
        # If no specific match, return a generic section
        return f"## {task}\n\nNo specific analysis results available for this task."
    
    @staticmethod
    def _generate_executive_summary(results: Dict[str, Any]) -> str:
        """Generate executive summary based on analysis results"""
        summary = []
        
        if "capacity_analysis" in results:
            capacity = results["capacity_analysis"]
            if "total_capacity" in capacity:
                summary.append(f"Total system capacity: {capacity['total_capacity']:,} beds")
            if "average_utilization" in capacity:
                summary.append(f"Average utilization: {capacity['average_utilization']:.1%}")
        
        if "facility_distribution" in results:
            dist = results["facility_distribution"]
            if "geographic_inequality" in dist:
                inequality = dist["geographic_inequality"]
                level = "High" if inequality > 0.4 else "Moderate" if inequality > 0.2 else "Low"
                summary.append(f"Geographic distribution inequality: {level} (Gini: {inequality:.2f})")
        
        if "recommendations" in results:
            high_priority = [r for r in results["recommendations"] if r.get("priority") == "High"]
            if high_priority:
                summary.append(f"{len(high_priority)} high-priority recommendations identified")
        
        if "validation" in results:
            completion_rate = results["validation"].get("completion_rate", 0)
            summary.append(f"Analysis completion rate: {completion_rate:.1%}")
        
        return " | ".join(summary) if summary else "No key metrics identified"
    
    @staticmethod
    def _format_validation_results(validation: Dict[str, Any]) -> str:
        """Format validation results"""
        response = f"**Analysis Completion Rate**: {validation['completion_rate']:.1%}\n\n"
        
        if validation["all_tasks_completed"]:
            response += "✅ All required tasks were completed successfully.\n\n"
        else:
            response += "⚠️ The following tasks were not completed:\n\n"
            for task in validation["missing_tasks"]:
                response += f"- {task}\n"
            response += "\n"
        
        return response
    
    # Legacy formatting methods for generic tasks
    @staticmethod
    def _format_facility_distribution(dist_results: Dict[str, Any]) -> str:
        """Format facility distribution analysis results"""
        response = ""
        
        if "total_facilities" in dist_results:
            response += f"**Total Facilities**: {dist_results['total_facilities']:,}\n\n"
        
        if "geographic_distribution" in dist_results:
            response += "### Geographic Distribution\n\n"
            response += "| Region | Facility Count |\n"
            response += "|--------|---------------|\n"
            for region, count in dist_results["geographic_distribution"].items():
                response += f"| {region} | {count} |\n"
            response += "\n"
            
            if "geographic_inequality" in dist_results:
                inequality = dist_results["geographic_inequality"]
                level = "High" if inequality > 0.4 else "Moderate" if inequality > 0.2 else "Low"
                response += f"**Geographic Inequality**: {level} (Gini coefficient: {inequality:.2f})\n\n"
        
        return response if response else "No facility distribution data available"
    
    @staticmethod
    def _format_resource_allocation(allocation_results: Dict[str, Any]) -> str:
        """Format resource allocation analysis results"""
        response = ""
        
        if "total_resources" in allocation_results:
            response += f"**Total Resources**: {allocation_results['total_resources']:,} units\n\n"
        
        if "allocation_by_facility" in allocation_results:
            response += "### Allocation by Facility Type\n\n"
            response += "| Facility Type | Resource Count | Percentage |\n"
            response += "|---------------|----------------|------------|\n"
            for ftype, count in allocation_results["allocation_by_facility"].items():
                percentage = count / allocation_results.get("total_resources", 1) * 100
                response += f"| {ftype} | {count:,} | {percentage:.1f}% |\n"
            response += "\n"
        
        return response if response else "No resource allocation data available"
    
    @staticmethod
    def _format_trends(trends_results: Dict[str, Any]) -> str:
        """Format trend analysis results"""
        response = ""
        
        if "summary" in trends_results:
            response += f"**Trend Summary**: {trends_results['summary']}\n\n"
        
        if "facility_growth" in trends_results:
            response += "### Facility Growth Trends\n\n"
            response += "| Year | Facility Count | Growth Rate |\n"
            response += "|------|----------------|-------------|\n"
            prev_count = None
            for year, count in trends_results["facility_growth"].items():
                if prev_count is not None:
                    growth_rate = (count - prev_count) / prev_count
                    response += f"| {year} | {count:,} | {growth_rate:.1%} |\n"
                else:
                    response += f"| {year} | {count:,} | N/A |\n"
                prev_count = count
            response += "\n"
        
        return response if response else "No trend data available"