"""
Flood Pulse Prediction System (FPPS) for cave systems.
Forecasts flood events and provides early warning.

Author: Samir Baladi
"""

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

logger = logging.getLogger(__name__)

@dataclass
class FloodForecast:
    """Represents a flood forecast."""
    time_to_peak: float  # hours
    peak_discharge: float  # m³/s
    peak_level: float  # m
    duration: float  # hours
    probability: float  # 0-1
    severity: str  # LOW, MODERATE, HIGH, EXTREME


class FloodForecasting:
    """
    Flood Pulse Prediction System (FPPS) for cave flood early warning.
    Based on the model from the research paper.
    """
    
    # Severity thresholds from paper
    SEVERITY_THRESHOLDS = {
        'LOW': 0.5,      # m³/s
        'MODERATE': 2.0, # m³/s
        'HIGH': 5.0,     # m³/s
        'EXTREME': 8.0   # m³/s
    }
    
    def __init__(self,
                 catchment_area: float = 12.5,  # km²
                 response_time: float = 6.0):    # hours
        """
        Initialize flood forecasting system.
        
        Args:
            catchment_area: Upstream catchment area (km²)
            response_time: Catchment response time (hours)
        """
        self.catchment_area = catchment_area
        self.response_time = response_time
        self.forecasts = []
        self.alerts = []
        logger.info(f"Initialized FloodForecasting with area={catchment_area}km²")

    def forecast_from_precipitation(self,
                                   precipitation: List[Dict],
                                   current_discharge: float,
                                   current_level: float,
                                   soil_saturation: float = 0.65) -> FloodForecast:
        """
        Forecast flood from precipitation data.
        
        Args:
            precipitation: List of precipitation events with amount and time
            current_discharge: Current discharge (m³/s)
            current_level: Current water level (m)
            soil_saturation: Soil saturation factor (0-1)
            
        Returns:
            FloodForecast object
        """
        # Calculate total precipitation
        total_precip = sum(p['amount'] for p in precipitation)
        
        # Calculate peak intensity
        peak_intensity = max(p['amount'] / 24 for p in precipitation)  # mm/hr
        
        # Runoff coefficient based on soil saturation
        runoff_coef = 0.2 + 0.6 * soil_saturation
        
        # Calculate peak discharge (rational method)
        # Q_peak = C * I * A / 360
        peak_discharge = (runoff_coef * peak_intensity * self.catchment_area) / 360
        
        # Add baseflow
        peak_discharge += current_discharge
        
        # Time to peak (depends on catchment and storm characteristics)
        time_to_peak = self.response_time + 0.5 * total_precip
        
        # Event duration
        duration = 24 + 2 * total_precip
        
        # Estimate peak level from discharge (simplified rating curve)
        # Assuming rectangular channel
        if current_discharge > 0:
            peak_level = current_level * (peak_discharge / current_discharge) ** 0.6
        else:
            peak_level = current_level + 1.0
        
        # Calculate probability based on data quality
        probability = min(0.95, 0.5 + 0.1 * len(precipitation))
        
        # Determine severity
        if peak_discharge >= self.SEVERITY_THRESHOLDS['EXTREME']:
            severity = "EXTREME"
        elif peak_discharge >= self.SEVERITY_THRESHOLDS['HIGH']:
            severity = "HIGH"
        elif peak_discharge >= self.SEVERITY_THRESHOLDS['MODERATE']:
            severity = "MODERATE"
        else:
            severity = "LOW"
        
        forecast = FloodForecast(
            time_to_peak=round(time_to_peak, 1),
            peak_discharge=round(peak_discharge, 2),
            peak_level=round(peak_level, 2),
            duration=round(duration, 1),
            probability=round(probability, 2),
            severity=severity
        )
        
        self.forecasts.append(forecast)
        
        # Generate alert if severe
        if severity in ['HIGH', 'EXTREME']:
            self._generate_alert(forecast)
        
        return forecast

    def forecast_from_rainfall_radar(self,
                                    radar_data: np.ndarray,
                                    timestamps: List[float],
                                    current_discharge: float,
                                    current_level: float) -> FloodForecast:
        """
        Forecast flood from rainfall radar data.
        
        Args:
            radar_data: Array of rainfall intensities (mm/hr)
            timestamps: Corresponding timestamps
            current_discharge: Current discharge (m³/s)
            current_level: Current water level (m)
            
        Returns:
            FloodForecast object
        """
        # Process radar data
        precipitation_events = []
        for i, intensity in enumerate(radar_data):
            if intensity > 0.1:  # Minimum threshold
                precipitation_events.append({
                    'time': timestamps[i],
                    'amount': intensity * 1,  # Assume 1 hour duration
                    'intensity': intensity
                })
        
        # Estimate soil saturation from recent rainfall
        recent_rainfall = sum(p['amount'] for p in precipitation_events[-24:])  # Last 24 hours
        soil_saturation = min(1.0, 0.3 + 0.02 * recent_rainfall)
        
        return self.forecast_from_precipitation(
            precipitation_events,
            current_discharge,
            current_level,
            soil_saturation
        )

    def update_forecast(self,
                       forecast_id: int,
                       new_data: Dict) -> FloodForecast:
        """
        Update an existing forecast with new data.
        
        Args:
            forecast_id: Index of forecast to update
            new_data: New data dictionary
            
        Returns:
            Updated FloodForecast
        """
        if forecast_id >= len(self.forecasts):
            raise ValueError(f"Forecast {forecast_id} not found")
        
        original = self.forecasts[forecast_id]
        
        # Update based on new data
        updated = FloodForecast(
            time_to_peak=max(0, original.time_to_peak - new_data.get('time_elapsed', 0)),
            peak_discharge=original.peak_discharge * (1 + new_data.get('intensity_factor', 0)),
            peak_level=original.peak_level * (1 + new_data.get('level_factor', 0)),
            duration=original.duration * (1 + new_data.get('duration_factor', 0)),
            probability=min(0.99, original.probability + new_data.get('confidence_increase', 0)),
            severity=original.severity
        )
        
        # Re-evaluate severity
        if updated.peak_discharge >= self.SEVERITY_THRESHOLDS['EXTREME']:
            updated.severity = "EXTREME"
        elif updated.peak_discharge >= self.SEVERITY_THRESHOLDS['HIGH']:
            updated.severity = "HIGH"
        elif updated.peak_discharge >= self.SEVERITY_THRESHOLDS['MODERATE']:
            updated.severity = "MODERATE"
        
        self.forecasts[forecast_id] = updated
        return updated

    def get_evacuation_timing(self,
                             forecast: FloodForecast,
                             exit_time: float,
                             safety_margin: float = 2.0) -> Dict[str, Union[str, float]]:
        """
        Calculate evacuation timing based on forecast.
        
        Args:
            forecast: Flood forecast
            exit_time: Time needed to exit (hours)
            safety_margin: Safety margin (hours)
            
        Returns:
            Dictionary with evacuation timing
        """
        total_time_needed = exit_time + safety_margin
        
        if forecast.time_to_peak <= 0:
            action = "FLOOD IN PROGRESS"
            evacuate_now = True
            time_available = 0
        elif forecast.time_to_peak <= total_time_needed:
            action = "EVACUATE IMMEDIATELY"
            evacuate_now = True
            time_available = forecast.time_to_peak
        elif forecast.time_to_peak <= total_time_needed * 2:
            action = "PREPARE TO EVACUATE"
            evacuate_now = False
            time_available = forecast.time_to_peak
        else:
            action = "MONITOR CONDITIONS"
            evacuate_now = False
            time_available = forecast.time_to_peak
        
        return {
            'action': action,
            'evacuate_now': evacuate_now,
            'time_available_hours': round(time_available, 1),
            'time_needed_hours': total_time_needed,
            'safety_margin_hours': safety_margin,
            'peak_in_hours': forecast.time_to_peak,
            'severity': forecast.severity
        }

    def calculate_flood_probability(self,
                                   current_discharge: float,
                                   historical_max: float,
                                   trend: float) -> float:
        """
        Calculate flood probability from current conditions.
        
        Args:
            current_discharge: Current discharge (m³/s)
            historical_max: Maximum historical discharge (m³/s)
            trend: Rate of change (m³/s per hour)
            
        Returns:
            Flood probability (0-1)
        """
        if historical_max <= 0:
            return 0.1
        
        # Base probability from current level
        base_prob = current_discharge / historical_max
        
        # Increase with rising trend
        trend_factor = min(0.3, trend * 0.1)
        
        # Combine
        probability = min(0.99, base_prob + trend_factor)
        
        return round(probability, 2)

    def _generate_alert(self, forecast: FloodForecast) -> None:
        """Generate flood alert."""
        alert = {
            'timestamp': datetime.now().timestamp(),
            'severity': forecast.severity,
            'peak_discharge': forecast.peak_discharge,
            'time_to_peak': forecast.time_to_peak,
            'probability': forecast.probability,
            'message': self._get_alert_message(forecast)
        }
        
        self.alerts.append(alert)
        logger.warning(f"FLOOD ALERT: {alert['message']}")

    def _get_alert_message(self, forecast: FloodForecast) -> str:
        """Get alert message based on severity."""
        if forecast.severity == "EXTREME":
            return (f"EXTREME FLOOD WARNING: {forecast.peak_discharge:.1f}m³/s peak "
                   f"in {forecast.time_to_peak:.1f} hours. IMMEDIATE EVACUATION REQUIRED.")
        elif forecast.severity == "HIGH":
            return (f"HIGH FLOOD WARNING: {forecast.peak_discharge:.1f}m³/s peak "
                   f"in {forecast.time_to_peak:.1f} hours. Prepare to evacuate.")
        else:
            return (f"FLOOD ADVISORY: {forecast.peak_discharge:.1f}m³/s peak "
                   f"in {forecast.time_to_peak:.1f} hours. Monitor conditions.")

    def get_active_alerts(self) -> List[Dict]:
        """Get active flood alerts."""
        current_time = datetime.now().timestamp()
        active = []
        
        for alert in self.alerts[-10:]:  # Last 10 alerts
            # Check if alert is still relevant (within last 24 hours)
            if current_time - alert['timestamp'] < 24 * 3600:
                active.append(alert)
        
        return active

    def summary(self) -> Dict[str, Union[int, float, str, List]]:
        """Get summary of flood forecasting."""
        if not self.forecasts:
            return {'status': 'No forecasts'}
        
        severities = [f.severity for f in self.forecasts[-10:]]
        
        return {
            'n_forecasts': len(self.forecasts),
            'n_alerts': len(self.alerts),
            'recent_severities': severities,
            'active_alerts': self.get_active_alerts(),
            'catchment_area': self.catchment_area,
            'response_time': self.response_time
        }


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Flood Forecasting (FPPS)")
    print("="*60)
    
    ff = FloodForecasting()
    
    # Precipitation events (from Vikos Gorge example)
    precipitation = [
        {'time': 0, 'amount': 25},
        {'time': 6, 'amount': 40},
        {'time': 12, 'amount': 15}
    ]
    
    # Generate forecast
    forecast = ff.forecast_from_precipitation(
        precipitation,
        current_discharge=0.15,
        current_level=0.45,
        soil_saturation=0.65
    )
    
    print(f"\nFlood Forecast:")
    print(f"  Severity: {forecast.severity}")
    print(f"  Peak Discharge: {forecast.peak_discharge} m³/s")
    print(f"  Time to Peak: {forecast.time_to_peak} hours")
    print(f"  Duration: {forecast.duration} hours")
    print(f"  Probability: {forecast.probability:.0%}")
    
    # Evacuation timing
    timing = ff.get_evacuation_timing(forecast, exit_time=2.0)
    print(f"\nEvacuation Timing:")
    print(f"  Action: {timing['action']}")
    print(f"  Time Available: {timing['time_available_hours']} hours")
    
    print("\n" + "="*60)
