"""
CDZ Risk Assessment - Evaluates hazards associated with Creative Dissolution Zones.
Provides risk scores, safety recommendations, and evacuation protocols.

Author: Samir Baladi
"""

import numpy as np
from typing import Dict, List, Tuple, Optional, Union
from dataclasses import dataclass
import logging

logger = logging.getLogger(__name__)

@dataclass
class RiskLevel:
    """Represents a risk level with associated actions."""
    level: str  # LOW, MODERATE, HIGH, EXTREME
    color: str
    score_range: Tuple[float, float]
    max_team_hours: int
    required_protocol: str
    evacuation_needed: bool


class CDZRiskAssessment:
    """
    Assesses risks associated with Creative Dissolution Zones.
    Provides quantitative risk scores and operational recommendations.
    """
    
    RISK_LEVELS = {
        'LOW': RiskLevel(
            level='LOW',
            color='🟢',
            score_range=(0, 25),
            max_team_hours=8,
            required_protocol='Standard Survey',
            evacuation_needed=False
        ),
        'MODERATE': RiskLevel(
            level='MODERATE',
            color='🟡',
            score_range=(25, 50),
            max_team_hours=6,
            required_protocol='Enhanced Monitoring',
            evacuation_needed=False
        ),
        'HIGH': RiskLevel(
            level='HIGH',
            color='🟠',
            score_range=(50, 75),
            max_team_hours=4,
            required_protocol='CDZ Protocol',
            evacuation_needed=False
        ),
        'VERY HIGH': RiskLevel(
            level='VERY HIGH',
            color='🔴',
            score_range=(75, 90),
            max_team_hours=2,
            required_protocol='Emergency Protocol',
            evacuation_needed=True
        ),
        'EXTREME': RiskLevel(
            level='EXTREME',
            color='⚫',
            score_range=(90, 100),
            max_team_hours=0,
            required_protocol='IMMEDIATE EVACUATION',
            evacuation_needed=True
        )
    }
    
    def __init__(self):
        """Initialize CDZ risk assessment."""
        self.assessments = []
        self.active_warnings = []
        logger.info("Initialized CDZRiskAssessment")

    def assess_risk(self,
                   cdz_proximity: float,
                   morphological_rate: float,
                   co2_level: float,
                   water_level_rise: float,
                   team_size: int = 3,
                   time_in_cave: float = 0) -> Dict[str, Union[str, float, List]]:
        """
        Assess overall risk based on multiple factors.
        
        Args:
            cdz_proximity: CDZ proximity score (0-1)
            morphological_rate: Rate of morphological change (m/season)
            co2_level: CO₂ concentration (ppm)
            water_level_rise: Rate of water level rise (m/hr)
            team_size: Number of team members
            time_in_cave: Time already spent in cave (hours)
            
        Returns:
            Dictionary with risk assessment
        """
        # Calculate individual risk components
        proximity_risk = self._assess_proximity_risk(cdz_proximity)
        morphological_risk = self._assess_morphological_risk(morphological_rate)
        co2_risk = self._assess_co2_risk(co2_level)
        water_risk = self._assess_water_risk(water_level_rise)
        team_risk = self._assess_team_risk(team_size, time_in_cave)
        
        # Combine risks with weights (from paper)
        weights = {
            'proximity': 0.35,
            'morphological': 0.25,
            'co2': 0.20,
            'water': 0.15,
            'team': 0.05
        }
        
        total_risk = (proximity_risk['score'] * weights['proximity'] +
                     morphological_risk['score'] * weights['morphological'] +
                     co2_risk['score'] * weights['co2'] +
                     water_risk['score'] * weights['water'] +
                     team_risk['score'] * weights['team'])
        
        # Determine risk level
        risk_level = self._get_risk_level(total_risk)
        
        # Calculate safe remaining time
        safe_time = self._calculate_safe_time(risk_level, time_in_cave)
        
        # Get required actions
        actions = self._get_required_actions(risk_level, {
            'proximity': proximity_risk,
            'morphological': morphological_risk,
            'co2': co2_risk,
            'water': water_risk,
            'team': team_risk
        })
        
        assessment = {
            'total_risk_score': round(total_risk, 1),
            'risk_level': risk_level.level,
            'color': risk_level.color,
            'max_team_hours': risk_level.max_team_hours,
            'time_remaining': safe_time,
            'evacuation_needed': risk_level.evacuation_needed,
            'required_protocol': risk_level.required_protocol,
            'components': {
                'proximity': proximity_risk,
                'morphological': morphological_risk,
                'co2': co2_risk,
                'water': water_risk,
                'team': team_risk
            },
            'actions': actions,
            'warnings': self._generate_warnings(proximity_risk, morphological_risk, 
                                                co2_risk, water_risk)
        }
        
        self.assessments.append(assessment)
        
        # Update active warnings
        if assessment['warnings']:
            self.active_warnings.extend(assessment['warnings'])
        
        return assessment

    def _assess_proximity_risk(self, proximity: float) -> Dict[str, Union[str, float]]:
        """Assess risk from CDZ proximity."""
        score = proximity * 100
        
        if proximity > 0.8:
            level = "EXTREME"
            description = "Inside active CDZ"
        elif proximity > 0.6:
            level = "VERY HIGH"
            description = "Very close to CDZ"
        elif proximity > 0.4:
            level = "HIGH"
            description = "Near CDZ boundary"
        elif proximity > 0.2:
            level = "MODERATE"
            description = "CDZ influence detected"
        else:
            level = "LOW"
            description = "Safe distance from CDZ"
        
        return {
            'score': round(score, 1),
            'level': level,
            'description': description,
            'proximity': proximity
        }

    def _assess_morphological_risk(self, rate: float) -> Dict[str, Union[str, float]]:
        """Assess risk from morphological change rate."""
        # Convert to score (0-100)
        score = min(100, rate * 200)  # 0.5 m/season = 100
        
        if rate > 0.4:
            level = "EXTREME"
            description = "Rapid morphological change"
        elif rate > 0.3:
            level = "VERY HIGH"
            description = "Active passage modification"
        elif rate > 0.2:
            level = "HIGH"
            description = "Significant change detected"
        elif rate > 0.1:
            level = "MODERATE"
            description = "Measurable change"
        else:
            level = "LOW"
            description = "Stable morphology"
        
        return {
            'score': round(score, 1),
            'level': level,
            'description': description,
            'rate': rate
        }

    def _assess_co2_risk(self, co2: float) -> Dict[str, Union[str, float]]:
        """Assess risk from CO₂ concentration."""
        # CO₂ thresholds from paper
        if co2 > 5000:
            score = 100
            level = "EXTREME"
            description = "IMMEDIATE DANGER - Life threatening"
        elif co2 > 3000:
            score = 80
            level = "VERY HIGH"
            description = "Severe respiratory hazard"
        elif co2 > 2000:
            score = 60
            level = "HIGH"
            description = "Significant CO₂ accumulation"
        elif co2 > 1000:
            score = 40
            level = "MODERATE"
            description = "Elevated CO₂ levels"
        else:
            score = 20
            level = "LOW"
            description = "Normal CO₂ levels"
        
        return {
            'score': round(score, 1),
            'level': level,
            'description': description,
            'co2_ppm': co2
        }

    def _assess_water_risk(self, rise_rate: float) -> Dict[str, Union[str, float]]:
        """Assess risk from water level rise."""
        # Convert to score (0-100)
        score = min(100, rise_rate * 200)  # 0.5 m/hr = 100
        
        if rise_rate > 0.4:
            level = "EXTREME"
            description = "Flash flood conditions"
        elif rise_rate > 0.3:
            level = "VERY HIGH"
            description = "Rapid water rise"
        elif rise_rate > 0.2:
            level = "HIGH"
            description = "Significant water rise"
        elif rise_rate > 0.1:
            level = "MODERATE"
            description = "Water level increasing"
        else:
            level = "LOW"
            description = "Stable water level"
        
        return {
            'score': round(score, 1),
            'level': level,
            'description': description,
            'rise_rate': rise_rate
        }

    def _assess_team_risk(self, team_size: int, time_in_cave: float) -> Dict[str, Union[str, float]]:
        """Assess risk from team factors."""
        # Base score from team size (smaller teams = higher risk)
        if team_size < 3:
            size_score = 60
            size_desc = "Minimum team size"
        elif team_size < 4:
            size_score = 40
            size_desc = "Adequate team size"
        else:
            size_score = 20
            size_desc = "Well-staffed team"
        
        # Time fatigue factor
        if time_in_cave > 8:
            fatigue_score = 40
            fatigue_desc = "Extended exposure"
        elif time_in_cave > 6:
            fatigue_score = 30
            fatigue_desc = "Moderate fatigue"
        elif time_in_cave > 4:
            fatigue_score = 20
            fatigue_desc = "Some fatigue"
        else:
            fatigue_score = 10
            fatigue_desc = "Fresh team"
        
        # Combined score
        score = (size_score + fatigue_score) / 2
        
        if score > 50:
            level = "HIGH"
        elif score > 35:
            level = "MODERATE"
        else:
            level = "LOW"
        
        return {
            'score': round(score, 1),
            'level': level,
            'team_size': team_size,
            'time_in_cave': time_in_cave,
            'size_description': size_desc,
            'fatigue_description': fatigue_desc
        }

    def _get_risk_level(self, score: float) -> RiskLevel:
        """Get risk level for a given score."""
        for level in self.RISK_LEVELS.values():
            if level.score_range[0] <= score < level.score_range[1]:
                return level
        return self.RISK_LEVELS['EXTREME']  # Default to extreme for scores >= 100

    def _calculate_safe_time(self, risk_level: RiskLevel, time_spent: float) -> float:
        """Calculate safe remaining time (hours)."""
        if risk_level.evacuation_needed:
            return 0
        
        max_time = risk_level.max_team_hours
        remaining = max(0, max_time - time_spent)
        
        return round(remaining, 1)

    def _get_required_actions(self, risk_level: RiskLevel, components: Dict) -> List[str]:
        """Get required actions based on risk level and components."""
        actions = []
        
        # Basic actions from risk level
        if risk_level.evacuation_needed:
            actions.append("IMMEDIATE EVACUATION - Exit cave now")
        else:
            actions.append(f"Maximum team time: {risk_level.max_team_hours} hours")
            actions.append(f"Required protocol: {risk_level.required_protocol}")
        
        # Component-specific actions
        if components['proximity']['level'] in ['EXTREME', 'VERY HIGH']:
            actions.append("Maintain maximum distance from CDZ center")
        
        if components['morphological']['level'] in ['EXTREME', 'VERY HIGH']:
            actions.append("Document all morphological changes")
            actions.append("Mark unstable areas with flagging")
        
        if components['co2']['level'] in ['EXTREME', 'VERY HIGH']:
            actions.append("Use respiratory protection")
            actions.append("Monitor CO₂ continuously")
        
        if components['water']['level'] in ['EXTREME', 'VERY HIGH']:
            actions.append("Move to higher ground immediately")
            actions.append("Monitor water level continuously")
        
        return actions

    def _generate_warnings(self, *risk_components) -> List[str]:
        """Generate warning messages for active risks."""
        warnings = []
        
        for component in risk_components:
            if component['level'] == 'EXTREME':
                warnings.append(f"EXTREME: {component['description']}")
            elif component['level'] == 'VERY HIGH':
                warnings.append(f"WARNING: {component['description']}")
        
        return warnings

    def get_evacuation_route(self,
                            current_position: Tuple[float, float, float],
                            exits: List[Tuple[float, float, float]],
                            risk_zones: List[Dict]) -> Dict[str, Union[List, float]]:
        """
        Get safest evacuation route.
        
        Args:
            current_position: (x, y, z) current location
            exits: List of exit coordinates
            risk_zones: List of risk zones with positions and risk levels
            
        Returns:
            Dictionary with evacuation route
        """
        # Simple implementation - find nearest exit avoiding high-risk zones
        # In practice, would use pathfinding algorithms
        
        # Filter out extreme risk zones
        safe_exits = []
        for exit_pos in exits:
            # Check if exit is in high-risk zone
            in_risk = False
            for zone in risk_zones:
                if zone.get('risk_level') in ['EXTREME', 'VERY HIGH']:
                    # Simple distance check
                    dist = np.sqrt(sum((exit_pos[j] - zone['position'][j])**2 for j in range(3)))
                    if dist < zone.get('radius', 50):
                        in_risk = True
                        break
            
            if not in_risk:
                # Calculate distance to exit
                dist = np.sqrt(sum((current_position[j] - exit_pos[j])**2 for j in range(3)))
                safe_exits.append({
                    'position': exit_pos,
                    'distance': dist,
                    'direction': self._get_direction(current_position, exit_pos)
                })
        
        if not safe_exits:
            return {
                'status': 'NO SAFE EXITS',
                'action': 'Take immediate shelter, call for rescue'
            }
        
        # Choose nearest safe exit
        nearest = min(safe_exits, key=lambda x: x['distance'])
        
        # Calculate estimated time (assuming 1 km/h travel speed in cave)
        travel_time = nearest['distance'] / 1000  # hours
        
        return {
            'recommended_exit': nearest['position'],
            'distance_m': round(nearest['distance'], 1),
            'direction': nearest['direction'],
            'estimated_time_hours': round(travel_time, 1),
            'alternative_exits': safe_exits[1:3] if len(safe_exits) > 1 else []
        }

    def _get_direction(self, pos1: Tuple, pos2: Tuple) -> str:
        """Get cardinal direction from pos1 to pos2."""
        dx = pos2[0] - pos1[0]
        dy = pos2[1] - pos1[1]
        
        if abs(dx) > abs(dy):
            return "East" if dx > 0 else "West"
        else:
            return "North" if dy > 0 else "South"

    def update_risk_status(self, new_measurements: Dict) -> List[str]:
        """
        Update risk status with new measurements.
        
        Args:
            new_measurements: Dictionary with new measurements
            
        Returns:
            List of new warnings
        """
        # Perform new assessment
        assessment = self.assess_risk(
            cdz_proximity=new_measurements.get('cdz_proximity', 0),
            morphological_rate=new_measurements.get('morphological_rate', 0),
            co2_level=new_measurements.get('co2', 400),
            water_level_rise=new_measurements.get('water_rise_rate', 0),
            team_size=new_measurements.get('team_size', 3),
            time_in_cave=new_measurements.get('time_in_cave', 0)
        )
        
        return assessment['warnings']

    def summary(self) -> Dict[str, Union[int, List, str]]:
        """Get summary of risk assessments."""
        if not self.assessments:
            return {'status': 'No assessments'}
        
        latest = self.assessments[-1]
        
        return {
            'n_assessments': len(self.assessments),
            'current_risk': latest['risk_level'],
            'current_color': latest['color'],
            'active_warnings': self.active_warnings[-5:] if self.active_warnings else [],
            'evacuation_status': 'ACTIVE' if latest['evacuation_needed'] else 'STANDING BY',
            'time_remaining': latest['time_remaining']
        }


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - CDZ Risk Assessment")
    print("="*60)
    
    ra = CDZRiskAssessment()
    
    # Assess risk for a scenario
    assessment = ra.assess_risk(
        cdz_proximity=0.75,
        morphological_rate=0.35,
        co2_level=2800,
        water_level_rise=0.15,
        team_size=3,
        time_in_cave=5.5
    )
    
    print(f"\nRisk Assessment:")
    print(f"  - Total Risk: {assessment['total_risk_score']}/100")
    print(f"  - Level: {assessment['color']} {assessment['risk_level']}")
    print(f"  - Protocol: {assessment['required_protocol']}")
    print(f"  - Time Remaining: {assessment['time_remaining']} hours")
    
    print(f"\nComponent Risks:")
    for comp, data in assessment['components'].items():
        print(f"  {comp}: {data['level']} ({data['score']})")
    
    print(f"\nActions Required:")
    for action in assessment['actions']:
        print(f"  • {action}")
    
    if assessment['warnings']:
        print(f"\nWarnings:")
        for warning in assessment['warnings']:
            print(f"  ⚠ {warning}")
    
    print("\n" + "="*60)
