import numpy as np
from typing import Dict, List, Optional
import json
from datetime import datetime, timedelta

class TreatmentPlanner:
    def __init__(self):
        self.treatment_protocols = self._load_treatment_protocols()
        self.intervention_strategies = self._load_intervention_strategies()
    
    def _load_treatment_protocols(self) -> Dict:
        """Load predefined treatment protocols"""
        return {
            'depression': {
                'psychotherapy': {
                    'primary': 'Cognitive Behavioral Therapy (CBT)',
                    'alternatives': [
                        'Interpersonal Therapy (IPT)',
                        'Behavioral Activation (BA)',
                        'Mindfulness-Based Cognitive Therapy (MBCT)'
                    ],
                    'duration': '12-16 weeks',
                    'frequency': 'Weekly'
                },
                'medication': {
                    'classes': [
                        'SSRIs',
                        'SNRIs',
                        'NDRIs',
                        'Atypical antidepressants'
                    ],
                    'duration': '6-12 months minimum',
                    'monitoring': 'Every 2-4 weeks initially'
                },
                'lifestyle': [
                    'Regular exercise (30 minutes daily)',
                    'Sleep hygiene improvement',
                    'Social engagement activities',
                    'Stress reduction techniques'
                ]
            },
            'anxiety': {
                'psychotherapy': {
                    'primary': 'Cognitive Behavioral Therapy (CBT)',
                    'alternatives': [
                        'Exposure Therapy',
                        'Acceptance and Commitment Therapy (ACT)',
                        'Dialectical Behavior Therapy (DBT)'
                    ],
                    'duration': '8-12 weeks',
                    'frequency': 'Weekly'
                },
                'medication': {
                    'classes': [
                        'SSRIs',
                        'SNRIs',
                        'Buspirone',
                        'Beta-blockers'
                    ],
                    'duration': 'As needed',
                    'monitoring': 'Every 2-4 weeks initially'
                },
                'lifestyle': [
                    'Relaxation techniques',
                    'Mindfulness meditation',
                    'Regular exercise',
                    'Stress management'
                ]
            },
            'ptsd': {
                'psychotherapy': {
                    'primary': 'Trauma-Focused CBT',
                    'alternatives': [
                        'EMDR',
                        'Prolonged Exposure Therapy',
                        'Cognitive Processing Therapy'
                    ],
                    'duration': '12-16 weeks',
                    'frequency': 'Weekly'
                },
                'medication': {
                    'classes': [
                        'SSRIs',
                        'SNRIs',
                        'Prazosin',
                        'Antipsychotics'
                    ],
                    'duration': '12 months minimum',
                    'monitoring': 'Every 2-4 weeks initially'
                },
                'lifestyle': [
                    'Stress management techniques',
                    'Sleep hygiene',
                    'Grounding exercises',
                    'Social support engagement'
                ]
            }
        }
    
    def _load_intervention_strategies(self) -> Dict:
        """Load intervention strategies based on severity and symptoms"""
        return {
            'mild': {
                'focus': 'Lifestyle modifications and psychoeducation',
                'monitoring': 'Monthly check-ins',
                'escalation_criteria': [
                    'Symptom worsening',
                    'Functional impairment',
                    'Lack of improvement after 4-6 weeks'
                ]
            },
            'moderate': {
                'focus': 'Psychotherapy with optional medication',
                'monitoring': 'Bi-weekly check-ins',
                'escalation_criteria': [
                    'Severe symptom exacerbation',
                    'Development of risk factors',
                    'Poor response to treatment'
                ]
            },
            'severe': {
                'focus': 'Intensive treatment with medication and therapy',
                'monitoring': 'Weekly check-ins',
                'escalation_criteria': [
                    'Crisis development',
                    'Safety concerns',
                    'Treatment resistance'
                ]
            }
        }
    
    def generate_plan(self, analysis_results: Dict) -> Dict:
        """Generate a comprehensive treatment plan based on analysis results"""
        try:
            # Extract relevant information from analysis
            conditions = analysis_results['condition_probabilities']
            severity = analysis_results['severity_assessment']
            
            # Generate treatment plan components
            treatment_plan = {
                'summary': self._generate_plan_summary(conditions, severity),
                'primary_interventions': self._select_primary_interventions(
                    conditions, severity
                ),
                'therapy_recommendations': self._generate_therapy_recommendations(
                    conditions, severity
                ),
                'medication_considerations': self._generate_medication_recommendations(
                    conditions, severity
                ),
                'lifestyle_modifications': self._generate_lifestyle_recommendations(
                    conditions
                ),
                'monitoring_plan': self._create_monitoring_plan(severity),
                'crisis_plan': self._create_crisis_plan(severity),
                'timeline': self._create_treatment_timeline(
                    conditions, severity
                )
            }
            
            return treatment_plan
            
        except Exception as e:
            return {
                'error': f'Error generating treatment plan: {str(e)}',
                'recommendations': self._generate_fallback_recommendations()
            }
    
    def _generate_plan_summary(
        self, conditions: Dict, severity: Dict
    ) -> Dict:
        """Generate a summary of the treatment plan"""
        primary_condition = max(conditions.items(), key=lambda x: x[1])[0]
        severity_level = self._determine_severity_level(severity['overall_severity'])
        
        return {
            'primary_condition': primary_condition,
            'severity_level': severity_level,
            'treatment_approach': self.intervention_strategies[severity_level]['focus'],
            'estimated_duration': self._estimate_treatment_duration(
                primary_condition, severity_level
            )
        }
    
    def _select_primary_interventions(
        self, conditions: Dict, severity: Dict
    ) -> List[str]:
        """Select primary interventions based on conditions and severity"""
        interventions = []
        severity_level = self._determine_severity_level(severity['overall_severity'])
        
        for condition, probability in conditions.items():
            if probability > 0.4:  # Consider conditions with significant probability
                protocol = self.treatment_protocols[condition]
                
                # Add primary therapy
                interventions.append(
                    f"Primary therapy: {protocol['psychotherapy']['primary']}"
                )
                
                # Add medication if moderate to severe
                if severity_level in ['moderate', 'severe']:
                    interventions.append(
                        f"Consider medication: {', '.join(protocol['medication']['classes'][:2])}"
                    )
        
        return interventions
    
    def _generate_therapy_recommendations(
        self, conditions: Dict, severity: Dict
    ) -> Dict:
        """Generate specific therapy recommendations"""
        therapy_plan = {}
        severity_level = self._determine_severity_level(severity['overall_severity'])
        
        for condition, probability in conditions.items():
            if probability > 0.3:
                protocol = self.treatment_protocols[condition]['psychotherapy']
                therapy_plan[condition] = {
                    'primary_therapy': protocol['primary'],
                    'alternatives': protocol['alternatives'][:2],
                    'frequency': protocol['frequency'],
                    'duration': protocol['duration']
                }
        
        return therapy_plan
    
    def _generate_medication_recommendations(
        self, conditions: Dict, severity: Dict
    ) -> Dict:
        """Generate medication recommendations"""
        medication_plan = {}
        severity_level = self._determine_severity_level(severity['overall_severity'])
        
        if severity_level in ['moderate', 'severe']:
            for condition, probability in conditions.items():
                if probability > 0.4:
                    protocol = self.treatment_protocols[condition]['medication']
                    medication_plan[condition] = {
                        'recommended_classes': protocol['classes'][:2],
                        'duration': protocol['duration'],
                        'monitoring': protocol['monitoring']
                    }
        
        return medication_plan
    
    def _generate_lifestyle_recommendations(self, conditions: Dict) -> List[str]:
        """Generate lifestyle modification recommendations"""
        recommendations = set()
        
        for condition, probability in conditions.items():
            if probability > 0.3:
                recommendations.update(
                    self.treatment_protocols[condition]['lifestyle']
                )
        
        return list(recommendations)
    
    def _create_monitoring_plan(self, severity: Dict) -> Dict:
        """Create a monitoring plan based on severity"""
        severity_level = self._determine_severity_level(severity['overall_severity'])
        strategy = self.intervention_strategies[severity_level]
        
        return {
            'frequency': strategy['monitoring'],
            'focus_areas': [
                'Symptom severity',
                'Treatment response',
                'Side effects',
                'Functional improvement'
            ],
            'escalation_criteria': strategy['escalation_criteria']
        }
    
    def _create_crisis_plan(self, severity: Dict) -> Dict:
        """Create a crisis intervention plan"""
        return {
            'warning_signs': [
                'Suicidal ideation',
                'Severe anxiety attacks',
                'Dissociative episodes',
                'Severe mood changes'
            ],
            'emergency_contacts': [
                'Primary therapist',
                'Crisis hotline',
                'Emergency services (911)',
                'Trusted support person'
            ],
            'immediate_actions': [
                'Contact emergency services if in immediate danger',
                'Use prescribed crisis medication if available',
                'Apply learned coping strategies',
                'Reach out to support system'
            ]
        }
    
    def _create_treatment_timeline(
        self, conditions: Dict, severity: Dict
    ) -> List[Dict]:
        """Create a timeline for treatment implementation"""
        timeline = []
        start_date = datetime.now()
        
        # Initial phase
        timeline.append({
            'phase': 'Initial Assessment and Stabilization',
            'duration': '1-2 weeks',
            'start_date': start_date.strftime('%Y-%m-%d'),
            'focus': 'Assessment and immediate interventions'
        })
        
        # Acute phase
        acute_start = start_date + timedelta(weeks=2)
        timeline.append({
            'phase': 'Acute Treatment',
            'duration': '8-12 weeks',
            'start_date': acute_start.strftime('%Y-%m-%d'),
            'focus': 'Primary interventions and symptom reduction'
        })
        
        # Continuation phase
        continuation_start = acute_start + timedelta(weeks=12)
        timeline.append({
            'phase': 'Continuation',
            'duration': '4-6 months',
            'start_date': continuation_start.strftime('%Y-%m-%d'),
            'focus': 'Maintaining improvements and preventing relapse'
        })
        
        return timeline
    
    def _determine_severity_level(self, severity_score: float) -> str:
        """Determine severity level from score"""
        if severity_score > 0.7:
            return 'severe'
        elif severity_score > 0.4:
            return 'moderate'
        else:
            return 'mild'
    
    def _estimate_treatment_duration(
        self, condition: str, severity_level: str
    ) -> str:
        """Estimate treatment duration based on condition and severity"""
        base_duration = {
            'mild': 3,
            'moderate': 6,
            'severe': 12
        }
        
        months = base_duration[severity_level]
        return f"{months}-{months+3} months"
    
    def _generate_fallback_recommendations(self) -> List[str]:
        """Generate basic recommendations when full plan generation fails"""
        return [
            "Seek professional mental health evaluation",
            "Consider psychotherapy",
            "Maintain regular sleep schedule",
            "Practice stress management techniques",
            "Engage in regular physical activity",
            "Build and maintain social support network"
        ]