"""
Flood Pulse Prediction System (FPPS)
Predicts discharge events and flood pulses in cave systems.

Equation: Q(t + Δt) = Q_base + A × sin(ω_LC × t) + Σ B_i × P(t - τ_i)

Author: Samir Baladi
"""

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

logger = logging.getLogger(__name__)

@dataclass
class FloodEvent:
    """Represents a flood event prediction."""
    peak_discharge: float  # Peak discharge (m³/s)
    time_to_peak: float  # Hours to peak
    duration: float  # Event duration (hours)
    water_level_rise: float  # Expected water level rise (m)
    confidence: float  # Prediction confidence (0-1)
    risk_level: str  # Risk level classification


class DischargePredictor:
    """
    Predicts discharge and flood events in cave systems.
    Implements the Flood Pulse Prediction System (FPPS) from the paper.
    """
    
    def __init__(self,
                 Q_base: float = 0.15,
                 A: float = 0.42,
                 omega_LC: float = 0.048,
                 precipitation_coefficients: Optional[List[float]] = None):
        """
        Initialize discharge predictor.
        
        Args:
            Q_base: Baseline flow (m³/s, default 0.15 from paper)
            A: Limit cycle amplitude (m³/s, default 0.42 from paper)
            omega_LC: Limit cycle frequency (yr⁻¹, default 0.048 from paper)
            precipitation_coefficients: Response coefficients for precipitation
        """
        self.Q_base = Q_base
        self.A = A
        self.omega_LC = omega_LC
        self.precipitation_coefficients = precipitation_coefficients or [0.3, 0.2, 0.1]
        self.discharge_history = []
        self.precipitation_history = []
        logger.info(f"Initialized DischargePredictor with Q_base={Q_base}m³/s")

    def add_discharge_measurement(self,
                                  timestamp: float,
                                  discharge: float,
                                  water_level: Optional[float] = None) -> None:
        """
        Add a discharge measurement.
        
        Args:
            timestamp: Time of measurement (hours or years)
            discharge: Measured discharge (m³/s)
            water_level: Optional water level (m)
        """
        self.discharge_history.append({
            'timestamp': timestamp,
            'discharge': discharge,
            'water_level': water_level
        })
        logger.debug(f"Added discharge measurement: {discharge:.2f}m³/s at t={timestamp}")

    def add_precipitation(self,
                         timestamp: float,
                         amount: float,
                         intensity: str = 'moderate') -> None:
        """
        Add precipitation data.
        
        Args:
            timestamp: Time of precipitation (hours)
            amount: Precipitation amount (mm)
            intensity: 'light', 'moderate', or 'heavy'
        """
        intensity_factor = {'light': 0.5, 'moderate': 1.0, 'heavy': 2.0}
        self.precipitation_history.append({
            'timestamp': timestamp,
            'amount': amount,
            'intensity': intensity,
            'factor': intensity_factor.get(intensity, 1.0)
        })
        logger.debug(f"Added precipitation: {amount}mm ({intensity}) at t={timestamp}")

    def predict_discharge(self,
                         current_time: float,
                         delta_t: float,
                         precipitation_forecast: Optional[List[Dict]] = None) -> Dict[str, float]:
        """
        Predict discharge at future time.
        
        Q(t + Δt) = Q_base + A × sin(ω_LC × t) + Σ B_i × P(t - τ_i)
        
        Args:
            current_time: Current time (hours)
            delta_t: Time ahead to predict (hours)
            precipitation_forecast: List of forecast precipitation events
            
        Returns:
            Dictionary with predicted discharge
        """
        future_time = current_time + delta_t
        
        # Base seasonal component
        t_years = future_time / 8760  # Convert hours to years
        seasonal = self.A * np.sin(self.omega_LC * t_years)
        
        # Precipitation response
        precip_response = 0.0
        if precipitation_forecast:
            for i, precip in enumerate(precipitation_forecast):
                if i < len(self.precipitation_coefficients):
                    lag = precip.get('lag', 6)  # Default lag 6 hours
                    if future_time - precip['timestamp'] <= lag:
                        B = self.precipitation_coefficients[i]
                        precip_response += B * precip['amount'] * precip.get('factor', 1.0)
        
        predicted = self.Q_base + seasonal + precip_response
        
        # Calculate confidence based on forecast horizon
        confidence = max(0.3, 1.0 - 0.1 * (delta_t / 24))  # Decreases with time
        
        return {
            'time': future_time,
            'discharge': round(predicted, 2),
            'seasonal_component': round(seasonal, 2),
            'precipitation_response': round(precip_response, 2),
            'confidence': round(confidence, 2)
        }

    def detect_flood_risk(self,
                          current_discharge: float,
                          forecast_hours: int = 48) -> Dict[str, Union[str, float]]:
        """
        Detect flood risk based on current and predicted discharge.
        
        Args:
            current_discharge: Current discharge (m³/s)
            forecast_hours: Hours to forecast ahead
            
        Returns:
            Dictionary with flood risk assessment
        """
        # Thresholds from paper
        thresholds = {
            'safe': 0.5,
            'caution': 2.0,
            'high_risk': 5.0,
            'critical': 5.0
        }
        
        # Current risk level
        if current_discharge < thresholds['safe']:
            current_risk = "SAFE"
            color = "🟢"
        elif current_discharge < thresholds['caution']:
            current_risk = "CAUTION"
            color = "🟡"
        elif current_discharge < thresholds['high_risk']:
            current_risk = "HIGH RISK"
            color = "🔴"
        else:
            current_risk = "CRITICAL"
            color = "⚫"
        
        # Predict future discharge (simple linear extrapolation)
        if len(self.discharge_history) >= 2:
            recent = self.discharge_history[-2:]
            rate = (recent[1]['discharge'] - recent[0]['discharge']) / (recent[1]['timestamp'] - recent[0]['timestamp'])
            predicted_peak = current_discharge + rate * forecast_hours
        else:
            predicted_peak = current_discharge * 1.5
        
        # Time to critical threshold
        if current_discharge < thresholds['critical'] and rate > 0:
            time_to_critical = (thresholds['critical'] - current_discharge) / rate if rate > 0 else float('inf')
        else:
            time_to_critical = 0
        
        return {
            'current_risk': current_risk,
            'color': color,
            'current_discharge': round(current_discharge, 2),
            'predicted_peak': round(predicted_peak, 2),
            'time_to_critical': round(time_to_critical, 1) if time_to_critical != float('inf') else '>48h',
            'threshold_safe': thresholds['safe'],
            'threshold_caution': thresholds['caution'],
            'threshold_high_risk': thresholds['high_risk'],
            'threshold_critical': thresholds['critical'],
            'evacuation_needed': current_discharge >= thresholds['high_risk'] or predicted_peak >= thresholds['critical']
        }

    def forecast_flood_event(self,
                            precipitation_events: List[Dict],
                            catchment_area: float = 12.5,
                            soil_saturation: float = 0.65) -> FloodEvent:
        """
        Forecast a specific flood event from precipitation.
        
        Args:
            precipitation_events: List of precipitation events
            catchment_area: Upstream catchment area (km²)
            soil_saturation: Pre-event soil saturation (0-1)
            
        Returns:
            FloodEvent object with prediction
        """
        # Calculate total precipitation
        total_precip = sum(p['amount'] for p in precipitation_events)
        
        # Runoff coefficient based on soil saturation
        runoff_coef = 0.2 + 0.6 * soil_saturation
        
        # Calculate peak discharge (rational method approximation)
        # Q_peak = C * I * A / 360 (simplified)
        peak_intensity = max(p['amount'] / 24 for p in precipitation_events)  # mm/hr
        peak_discharge = runoff_coef * peak_intensity * catchment_area / 360  # m³/s
        
        # Add baseflow
        peak_discharge += self.Q_base
        
        # Time to peak (depends on catchment characteristics)
        time_to_peak = 6 + 2 * total_precip  # hours
        
        # Event duration
        duration = 24 + 3 * total_precip  # hours
        
        # Water level rise (simplified)
        water_level_rise = 0.5 * peak_discharge / 5  # m
        
        # Confidence based on data quality
        confidence = min(0.95, 0.5 + 0.3 * len(precipitation_events) / 5)
        
        # Risk level
        if peak_discharge < 0.5:
            risk = "LOW"
        elif peak_discharge < 2.0:
            risk = "MODERATE"
        elif peak_discharge < 5.0:
            risk = "HIGH"
        else:
            risk = "EXTREME"
        
        return FloodEvent(
            peak_discharge=round(peak_discharge, 2),
            time_to_peak=round(time_to_peak, 1),
            duration=round(duration, 1),
            water_level_rise=round(water_level_rise, 2),
            confidence=round(confidence, 2),
            risk_level=risk
        )

    def calculate_recurrence_interval(self,
                                     discharge: float,
                                     historical_data: Optional[List[float]] = None) -> Dict[str, float]:
        """
        Calculate recurrence interval for a given discharge.
        
        Args:
            discharge: Discharge value (m³/s)
            historical_data: Optional historical discharge data
            
        Returns:
            Dictionary with recurrence statistics
        """
        if historical_data is None and self.discharge_history:
            historical_data = [d['discharge'] for d in self.discharge_history]
        
        if not historical_data or len(historical_data) < 10:
            return {
                'recurrence_years': 5.0,
                'probability_annual': 0.2,
                'confidence': 0.3,
                'message': 'Insufficient historical data'
            }
        
        # Fit to Gumbel distribution (extreme value)
        sorted_data = np.sort(historical_data)
        
        # Empirical probability
        n = len(sorted_data)
        exceedance = np.sum(sorted_data >= discharge) / n
        
        if exceedance > 0:
            recurrence = 1 / exceedance
        else:
            recurrence = n + 1  # Beyond observed range
        
        return {
            'recurrence_years': round(recurrence, 1),
            'probability_annual': round(1/recurrence, 3) if recurrence > 0 else 0,
            'exceedance_probability': round(exceedance, 3),
            'n_events': n,
            'confidence': round(min(0.9, 0.3 + 0.02 * n), 2)
        }

    def get_evacuation_timing(self,
                             current_discharge: float,
                             discharge_rate: float) -> Dict[str, Union[float, str]]:
        """
        Calculate evacuation timing based on discharge.
        
        Args:
            current_discharge: Current discharge (m³/s)
            discharge_rate: Rate of change (m³/s per hour)
            
        Returns:
            Dictionary with evacuation timing
        """
        thresholds = {
            'prepare': 2.0,
            'evacuate': 5.0
        }
        
        if current_discharge >= thresholds['evacuate']:
            return {
                'action': 'IMMEDIATE EVACUATION',
                'time_available': 0,
                'message': 'Critical level reached - exit now'
            }
        
        if discharge_rate <= 0:
            return {
                'action': 'MONITOR',
                'time_available': float('inf'),
                'message': 'Discharge stable or decreasing'
            }
        
        # Time to prepare threshold
        time_to_prepare = (thresholds['prepare'] - current_discharge) / discharge_rate if current_discharge < thresholds['prepare'] else 0
        
        # Time to evacuate threshold
        time_to_evacuate = (thresholds['evacuate'] - current_discharge) / discharge_rate
        
        if time_to_evacuate < 2:
            action = "IMMEDIATE EVACUATION"
            message = f"Only {time_to_evacuate:.1f} hours until critical"
        elif time_to_evacuate < 6:
            action = "PREPARE TO EVACUATE"
            message = f"{time_to_evacuate:.1f} hours until critical"
        elif time_to_prepare < 2:
            action = "CAUTION - MONITOR CLOSELY"
            message = f"{time_to_prepare:.1f} hours until prepare threshold"
        else:
            action = "MONITOR"
            message = f"Safe for now, {time_to_evacuate:.1f} hours until critical"
        
        return {
            'action': action,
            'time_to_prepare': round(time_to_prepare, 1),
            'time_to_evacuate': round(time_to_evacuate, 1),
            'current_discharge': round(current_discharge, 2),
            'discharge_rate': round(discharge_rate, 3),
            'message': message
        }

    def summary(self) -> Dict[str, Union[float, int, str]]:
        """Get summary of discharge predictions."""
        return {
            'baseline_flow': self.Q_base,
            'seasonal_amplitude': self.A,
            'seasonal_period': round(2 * np.pi / self.omega_LC / 8760, 1),  # in hours
            'n_measurements': len(self.discharge_history),
            'n_precipitation_events': len(self.precipitation_history),
            'precipitation_coefficients': self.precipitation_coefficients
        }


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Discharge Prediction (FPPS)")
    print("="*60)
    
    predictor = DischargePredictor()
    
    # Add some historical data
    predictor.add_discharge_measurement(0, 0.15)
    predictor.add_discharge_measurement(24, 0.22)
    predictor.add_discharge_measurement(48, 0.35)
    
    # Predict discharge
    prediction = predictor.predict_discharge(48, 24)
    print(f"\nDischarge Prediction (24h ahead):")
    print(f"  - Predicted: {prediction['discharge']} m³/s")
    print(f"  - Confidence: {prediction['confidence']:.0%}")
    
    # Detect flood risk
    risk = predictor.detect_flood_risk(0.35)
    print(f"\nFlood Risk Assessment:")
    print(f"  - Current Risk: {risk['color']} {risk['current_risk']}")
    print(f"  - Evacuation Needed: {risk['evacuation_needed']}")
    
    # Forecast flood event
    precip_events = [
        {'timestamp': 0, 'amount': 25, 'intensity': 'moderate'},
        {'timestamp': 6, 'amount': 40, 'intensity': 'heavy'},
        {'timestamp': 12, 'amount': 15, 'intensity': 'light'}
    ]
    
    flood = predictor.forecast_flood_event(precip_events)
    print(f"\nFlood Event Forecast:")
    print(f"  - Peak Discharge: {flood.peak_discharge} m³/s")
    print(f"  - Time to Peak: {flood.time_to_peak} hours")
    print(f"  - Risk Level: {flood.risk_level}")
    print(f"  - Confidence: {flood.confidence:.0%}")
    
    # Evacuation timing
    timing = predictor.get_evacuation_timing(0.35, 0.05)
    print(f"\nEvacuation Timing:")
    print(f"  - Action: {timing['action']}")
    print(f"  - {timing['message']}")
    
    print("\n" + "="*60)
