"""
Stability envelope calculations for cave passage systems.
Implements APSEP (Adaptive Passage Stability Envelope Prediction) framework.

Author: Samir Baladi
"""

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

logger = logging.getLogger(__name__)

@dataclass
class PassageState:
    """Represents the current state of a cave passage."""
    V: float  # Vertical passage dimension (m)
    L: float  # Lateral conduit width (m)
    F: float  # Flowstone thickness (mm)
    Q: Optional[float] = None  # Water discharge (m³/s)
    CO2: Optional[float] = None  # CO₂ concentration (ppm)
    T: Optional[float] = None  # Temperature (°C)
    RH: Optional[float] = None  # Relative humidity (%)


@dataclass
class StabilityEnvelope:
    """Represents the stability envelope for a passage state."""
    score: float  # Stability score (0-100)
    level: str  # STABLE, CAUTION, UNSTABLE, HIGH RISK, CRITICAL
    color: str  # 🟢, 🟡, 🟠, 🔴, ⚫
    alert: str  # None, Advisory, Warning, Alert, Emergency
    action: str  # Recommended action
    max_hours: str  # Maximum team hours
    p_cdz: float  # CDZ proximity (0-1)
    safe_until: Optional[float] = None  # Predicted safe time (hours)


class StabilityEnvelopeCalculator:
    """
    Calculates stability envelopes for cave passages based on VLF dynamics.
    Implements the APSEP (Adaptive Passage Stability Envelope Prediction) framework.
    """
    
    def __init__(self, 
                 threshold_stable: float = 80.0,
                 threshold_caution: float = 60.0,
                 threshold_unstable: float = 40.0,
                 threshold_high_risk: float = 20.0):
        """
        Initialize the stability envelope calculator.
        
        Args:
            threshold_stable: Score threshold for STABLE (80-100)
            threshold_caution: Score threshold for CAUTION (60-79)
            threshold_unstable: Score threshold for UNSTABLE (40-59)
            threshold_high_risk: Score threshold for HIGH RISK (20-39)
        """
        self.thresholds = {
            'STABLE': threshold_stable,
            'CAUTION': threshold_caution,
            'UNSTABLE': threshold_unstable,
            'HIGH_RISK': threshold_high_risk,
            'CRITICAL': 0.0
        }
        logger.info("Initialized StabilityEnvelopeCalculator")

    def calculate_envelope(self, 
                          state: PassageState,
                          historical_data: Optional[np.ndarray] = None,
                          forecast_hours: int = 48) -> StabilityEnvelope:
        """
        Calculate stability envelope for a given passage state.
        
        Args:
            state: Current passage state
            historical_data: Optional historical time series
            forecast_hours: Hours to forecast ahead
            
        Returns:
            StabilityEnvelope with assessment
        """
        logger.info(f"Calculating stability envelope for state V={state.V}, L={state.L}, F={state.F}")
        
        # Calculate base stability score
        base_score = self._calculate_base_score(state)
        
        # Apply hydrological adjustment if Q available
        if state.Q is not None:
            base_score = self._apply_hydrological_adjustment(base_score, state.Q)
        
        # Apply CO₂ adjustment if available
        if state.CO2 is not None:
            base_score = self._apply_co2_adjustment(base_score, state.CO2)
        
        # Calculate CDZ proximity
        p_cdz = self._calculate_cdz_proximity(state)
        
        # Adjust score based on CDZ proximity
        adjusted_score = base_score * (1.0 - 0.3 * p_cdz)
        
        # Forecast stability horizon
        safe_until = self._forecast_stability_horizon(state, adjusted_score, forecast_hours)
        
        # Determine stability level
        level, color, alert, action, max_hours = self._classify_stability(adjusted_score)
        
        return StabilityEnvelope(
            score=round(adjusted_score, 1),
            level=level,
            color=color,
            alert=alert,
            action=action,
            max_hours=max_hours,
            p_cdz=round(p_cdz, 3),
            safe_until=safe_until
        )

    def _calculate_base_score(self, state: PassageState) -> float:
        """
        Calculate base stability score from V, L, F values.
        """
        # Ideal values from paper
        V_ideal = 2.5  # m
        L_ideal = 3.0  # m
        F_ideal = 0.08  # mm/yr
        
        # Calculate deviations
        V_dev = abs(state.V - V_ideal) / V_ideal
        L_dev = abs(state.L - L_ideal) / L_ideal
        F_dev = abs(state.F - F_ideal) / F_ideal
        
        # Weighted combination (weights from coupling coefficients)
        weights = {'V': 0.4, 'L': 0.35, 'F': 0.25}
        total_dev = (weights['V'] * V_dev + 
                    weights['L'] * L_dev + 
                    weights['F'] * F_dev)
        
        # Convert to score (100 - scaled deviation)
        score = max(0, 100 * (1 - min(1, total_dev)))
        
        return score

    def _apply_hydrological_adjustment(self, score: float, Q: float) -> float:
        """
        Adjust stability score based on discharge conditions.
        """
        # Normal flow: 0.15 m³/s (from paper)
        Q_normal = 0.15
        
        if Q < 0.5:  # Safe conditions
            return score
        elif Q < 2.0:  # Caution
            return score * 0.9
        elif Q < 5.0:  # High risk
            return score * 0.7
        else:  # Critical
            return score * 0.4

    def _apply_co2_adjustment(self, score: float, CO2: float) -> float:
        """
        Adjust stability score based on CO₂ concentration.
        """
        # Normal CO2: 400-1000 ppm
        if CO2 < 1000:
            return score
        elif CO2 < 2500:
            return score * 0.95
        elif CO2 < 5000:
            return score * 0.85
        elif CO2 < 8500:
            return score * 0.7
        else:
            return score * 0.5

    def _calculate_cdz_proximity(self, state: PassageState) -> float:
        """
        Calculate proximity to Creative Dissolution Zone (0-1).
        """
        # Check if near phase transition boundaries
        V_gradient = self._estimate_vertical_gradient(state.V)
        L_gradient = self._estimate_lateral_gradient(state.L)
        
        # CDZ criteria from paper
        if (V_gradient > 0.5 or L_gradient > 0.5):
            return min(1.0, (V_gradient + L_gradient) / 2)
        else:
            return 0.0

    def _estimate_vertical_gradient(self, V: float) -> float:
        """
        Estimate vertical passage gradient (rate of change).
        Simplified - in practice use historical data.
        """
        # Check if V is in transition zone (phreatic-vadose interface ~2-3m)
        if 2.0 <= V <= 3.0:
            return 0.3  # Moderate gradient
        elif V < 1.0 or V > 4.0:
            return 0.1  # Low gradient
        else:
            return 0.2

    def _estimate_lateral_gradient(self, L: float) -> float:
        """
        Estimate lateral conduit gradient.
        """
        # Check if L is in transition zone
        if 2.5 <= L <= 3.5:
            return 0.3
        elif L < 1.5 or L > 4.5:
            return 0.1
        else:
            return 0.2

    def _forecast_stability_horizon(self, 
                                    state: PassageState, 
                                    score: float, 
                                    max_hours: int) -> Optional[float]:
        """
        Forecast how long current stability will last.
        
        Returns:
            Hours until stability degrades (or None if stable)
        """
        if score >= 80:
            return None  # Indefinitely stable
        
        # Simple exponential decay model
        # In practice, use full VLF dynamics simulation
        decay_rate = 0.05  # per hour
        hours_to_threshold = np.log(self.thresholds['CAUTION'] / score) / (-decay_rate)
        
        return min(max_hours, max(0, hours_to_threshold))

    def _classify_stability(self, score: float) -> Tuple[str, str, str, str, str]:
        """
        Classify stability level based on score.
        
        Returns:
            Tuple of (level, color, alert, action, max_hours)
        """
        if score >= 80:
            return ("STABLE", "🟢", "None", 
                   "Standard Survey Mode", "Unrestricted")
        elif score >= 60:
            return ("CAUTION", "🟡", "Advisory",
                   "Enhanced Monitoring", "8 hours/shift")
        elif score >= 40:
            return ("UNSTABLE", "🟠", "Warning",
                   "CDZ Protocol Active", "6 hours/shift")
        elif score >= 20:
            return ("HIGH RISK", "🔴", "Alert",
                   "Evacuation Preparation", "4 hours/shift")
        else:
            return ("CRITICAL", "⚫", "Emergency",
                   "Immediate Evacuation", "0 hours (exit)")

    def batch_analyze(self, 
                     states: List[PassageState],
                     timestamps: Optional[np.ndarray] = None) -> List[StabilityEnvelope]:
        """
        Analyze multiple passage states.
        
        Args:
            states: List of passage states
            timestamps: Optional timestamps for each state
            
        Returns:
            List of stability envelopes
        """
        logger.info(f"Batch analyzing {len(states)} states")
        results = []
        
        for i, state in enumerate(states):
            envelope = self.calculate_envelope(state)
            results.append(envelope)
            
            if i % 10 == 0:
                logger.info(f"Processed {i}/{len(states)} states")
        
        return results

    def get_recommendation(self, envelope: StabilityEnvelope) -> Dict[str, str]:
        """
        Get detailed recommendation based on stability envelope.
        """
        recommendations = {
            'STABLE': {
                'summary': 'Safe for standard survey operations',
                'equipment': 'Standard survey gear sufficient',
                'team': '2-3 members, normal rotation',
                'documentation': 'Normal interval (every 5m)',
                'monitoring': 'Standard sensor deployment'
            },
            'CAUTION': {
                'summary': 'Exercise caution, monitor conditions',
                'equipment': 'Additional sensors recommended',
                'team': '3-4 members with sensor specialist',
                'documentation': 'Increased interval (every 2m)',
                'monitoring': 'Continuous CO2 and water level'
            },
            'UNSTABLE': {
                'summary': 'Active CDZ conditions, enhanced protocols',
                'equipment': 'Full sensor suite required',
                'team': '5-6 members with safety backup',
                'documentation': 'Maximum detail (every 0.5m)',
                'monitoring': 'Real-time continuous monitoring'
            },
            'HIGH RISK': {
                'summary': 'Prepare for evacuation',
                'equipment': 'Emergency gear ready',
                'team': 'Prepare to exit',
                'documentation': 'Only critical measurements',
                'monitoring': 'Evacuation monitoring only'
            },
            'CRITICAL': {
                'summary': 'IMMEDIATE EVACUATION REQUIRED',
                'equipment': 'Abandon non-essential gear',
                'team': 'Exit immediately',
                'documentation': 'None - focus on safety',
                'monitoring': 'Disable - exit immediately'
            }
        }
        
        return recommendations.get(envelope.level, {})


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Stability Envelope Analysis")
    print("="*60)
    
    calculator = StabilityEnvelopeCalculator()
    
    # Example 1: Stable passage (Vikos Gorge normal conditions)
    state1 = PassageState(V=2.8, L=3.2, F=0.087, Q=0.15, CO2=850)
    envelope1 = calculator.calculate_envelope(state1)
    
    print(f"\n1. Stable Passage Assessment:")
    print(f"   - Score: {envelope1.score}/100")
    print(f"   - Level: {envelope1.color} {envelope1.level}")
    print(f"   - CDZ Proximity: {envelope1.p_cdz:.1%}")
    print(f"   - Action: {envelope1.action}")
    
    # Example 2: Unstable passage (CDZ active)
    state2 = PassageState(V=3.8, L=4.5, F=0.12, Q=2.5, CO2=3500)
    envelope2 = calculator.calculate_envelope(state2)
    
    print(f"\n2. Unstable Passage Assessment:")
    print(f"   - Score: {envelope2.score}/100")
    print(f"   - Level: {envelope2.color} {envelope2.level}")
    print(f"   - CDZ Proximity: {envelope2.p_cdz:.1%}")
    print(f"   - Action: {envelope2.action}")
    if envelope2.safe_until:
        print(f"   - Safe Until: {envelope2.safe_until:.1f} hours")
    
    # Example 3: Critical passage (flood conditions)
    state3 = PassageState(V=1.2, L=6.8, F=0.15, Q=6.5, CO2=7200)
    envelope3 = calculator.calculate_envelope(state3)
    
    print(f"\n3. Critical Passage Assessment:")
    print(f"   - Score: {envelope3.score}/100")
    print(f"   - Level: {envelope3.color} {envelope3.level}")
    print(f"   - Action: {envelope3.action}")
    
    # Get recommendations
    print(f"\nRecommendations for {envelope2.level}:")
    rec = calculator.get_recommendation(envelope2)
    for key, value in rec.items():
        print(f"   - {key}: {value}")
    
    print("\n" + "="*60)
