"""
Adaptive Passage Stability Envelope Prediction (APSEP) framework.
Predicts stability envelopes for cave passages based on VLF dynamics.

Author: Samir Baladi
"""

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

logger = logging.getLogger(__name__)

@dataclass
class StabilityEnvelope:
    """Represents a stability envelope prediction."""
    time: float  # Prediction time
    v_min: float  # Minimum V dimension (m)
    v_max: float  # Maximum V dimension (m)
    l_min: float  # Minimum L dimension (m)
    l_max: float  # Maximum L dimension (m)
    f_min: float  # Minimum F thickness (mm)
    f_max: float  # Maximum F thickness (mm)
    probability: float  # Probability of stability (0-1)
    confidence: float  # Prediction confidence (0-1)


class StabilityEnvelopePredictor:
    """
    Predicts stability envelopes using the APSEP framework.
    Based on the coupled VLF oscillator model from the paper.
    """
    
    def __init__(self,
                 mu_v: float = 0.15,
                 mu_l: float = 0.10,
                 mu_f: float = 0.05,
                 k_vf: float = 0.23,
                 k_lf: float = -0.18,
                 k_fv: float = 0.31,
                 k_fl: float = -0.25):
        """
        Initialize stability envelope predictor.
        
        Args:
            mu_v, mu_l, mu_f: Non-linearity parameters
            k_vf, k_lf, k_fv, k_fl: Coupling coefficients from paper
        """
        self.mu_v = mu_v
        self.mu_l = mu_l
        self.mu_f = mu_f
        self.k_vf = k_vf
        self.k_lf = k_lf
        self.k_fv = k_fv
        self.k_fl = k_fl
        self.predictions = []
        logger.info("Initialized StabilityEnvelopePredictor")

    def predict_envelope(self,
                        current_state: Tuple[float, float, float],
                        time_horizon: float,
                        n_points: int = 100,
                        confidence_level: float = 0.95) -> StabilityEnvelope:
        """
        Predict stability envelope for a given time horizon.
        
        Args:
            current_state: Current (V, L, F) values
            time_horizon: Prediction horizon (years)
            n_points: Number of trajectory points
            confidence_level: Confidence level for envelope
            
        Returns:
            StabilityEnvelope prediction
        """
        V0, L0, F0 = current_state
        
        # Simulate multiple trajectories with noise
        n_trajectories = 100
        all_V = []
        all_L = []
        all_F = []
        
        for _ in range(n_trajectories):
            # Add small random perturbations to initial conditions
            V_noise = V0 * (1 + np.random.randn() * 0.05)
            L_noise = L0 * (1 + np.random.randn() * 0.05)
            F_noise = F0 * (1 + np.random.randn() * 0.05)
            
            # Simulate trajectory
            t, V, L, F = self._simulate_trajectory(
                (V_noise, L_noise, F_noise),
                time_horizon,
                n_points
            )
            
            all_V.append(V)
            all_L.append(L)
            all_F.append(F)
        
        # Convert to arrays
        all_V = np.array(all_V)
        all_L = np.array(all_L)
        all_F = np.array(all_F)
        
        # Calculate envelope at each time point
        V_min = np.percentile(all_V, (1 - confidence_level) * 50, axis=0)
        V_max = np.percentile(all_V, 100 - (1 - confidence_level) * 50, axis=0)
        L_min = np.percentile(all_L, (1 - confidence_level) * 50, axis=0)
        L_max = np.percentile(all_L, 100 - (1 - confidence_level) * 50, axis=0)
        F_min = np.percentile(all_F, (1 - confidence_level) * 50, axis=0)
        F_max = np.percentile(all_F, 100 - (1 - confidence_level) * 50, axis=0)
        
        # Calculate stability probability
        # Probability that trajectory stays within bounds
        bounds = self._calculate_bounds(current_state)
        stable_count = 0
        for i in range(n_trajectories):
            V_traj = all_V[i]
            L_traj = all_L[i]
            F_traj = all_F[i]
            
            if (np.all(V_traj >= bounds['V_min']) and 
                np.all(V_traj <= bounds['V_max']) and
                np.all(L_traj >= bounds['L_min']) and 
                np.all(L_traj <= bounds['L_max']) and
                np.all(F_traj >= bounds['F_min']) and 
                np.all(F_traj <= bounds['F_max'])):
                stable_count += 1
        
        probability = stable_count / n_trajectories
        
        # Create envelope at final time
        envelope = StabilityEnvelope(
            time=time_horizon,
            v_min=float(V_min[-1]),
            v_max=float(V_max[-1]),
            l_min=float(L_min[-1]),
            l_max=float(L_max[-1]),
            f_min=float(F_min[-1]),
            f_max=float(F_max[-1]),
            probability=probability,
            confidence=confidence_level
        )
        
        self.predictions.append(envelope)
        return envelope

    def _simulate_trajectory(self,
                            initial_state: Tuple[float, float, float],
                            time_horizon: float,
                            n_points: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """
        Simulate a single VLF trajectory.
        """
        V0, L0, F0 = initial_state
        
        def derivatives(t, state):
            V, L, F = state
            dVdt = L - self.mu_v * (V**2 - 1) * V + self.k_vf * F
            dLdt = -V + self.mu_l * (1 - L**2) * L + self.k_lf * F
            dFdt = -self.k_fv * V - self.k_fl * L + self.mu_f * (1 - F**2) * F
            return [dVdt, dLdt, dFdt]
        
        t_span = (0, time_horizon)
        t_eval = np.linspace(0, time_horizon, n_points)
        
        solution = integrate.solve_ivp(
            derivatives,
            t_span,
            [V0, L0, F0],
            t_eval=t_eval,
            method='RK45',
            rtol=1e-6
        )
        
        return t_eval, solution.y[0], solution.y[1], solution.y[2]

    def _calculate_bounds(self, state: Tuple[float, float, float]) -> Dict[str, float]:
        """
        Calculate stability bounds for a state.
        """
        V, L, F = state
        
        # Bounds based on typical stable ranges from paper
        return {
            'V_min': max(0.5, V * 0.7),
            'V_max': min(10, V * 1.3),
            'L_min': max(0.5, L * 0.7),
            'L_max': min(15, L * 1.3),
            'F_min': max(0.01, F * 0.5),
            'F_max': min(1.0, F * 1.5)
        }

    def predict_safe_time(self,
                         current_state: Tuple[float, float, float],
                         threshold: float = 0.8) -> Dict[str, float]:
        """
        Predict safe exploration time.
        
        Args:
            current_state: Current (V, L, F) values
            threshold: Minimum stability probability threshold
            
        Returns:
            Dictionary with safe time predictions
        """
        safe_time = 0
        time_step = 0.1  # years
        
        while safe_time < 5:  # Max 5 years prediction
            envelope = self.predict_envelope(
                current_state,
                safe_time + time_step,
                confidence_level=0.95
            )
            
            if envelope.probability < threshold:
                break
            
            safe_time += time_step
        
        return {
            'safe_time_years': round(safe_time, 1),
            'safe_time_hours': round(safe_time * 365.25 * 24),
            'confidence': envelope.confidence,
            'threshold': threshold
        }

    def get_stability_trend(self,
                           history: List[Tuple[float, float, float]],
                           times: List[float]) -> Dict[str, Union[str, float]]:
        """
        Analyze stability trend from historical data.
        
        Args:
            history: List of (V, L, F) states
            times: Corresponding times
            
        Returns:
            Dictionary with trend analysis
        """
        if len(history) < 3:
            return {'status': 'Insufficient data'}
        
        # Calculate stability scores over time
        scores = []
        for state in history:
            envelope = self.predict_envelope(state, 1.0, confidence_level=0.8)
            scores.append(envelope.probability)
        
        # Fit trend
        coeffs = np.polyfit(times, scores, 1)
        slope = coeffs[0]
        
        if slope > 0.05:
            trend = "IMPROVING"
            color = "🟢"
        elif slope < -0.05:
            trend = "DEGRADING"
            color = "🔴"
        else:
            trend = "STABLE"
            color = "🟡"
        
        return {
            'trend': trend,
            'color': color,
            'slope_per_year': round(slope, 3),
            'current_stability': round(scores[-1], 2),
            'initial_stability': round(scores[0], 2),
            'change': round(scores[-1] - scores[0], 2)
        }

    def summary(self) -> Dict[str, Union[int, float, str]]:
        """Get summary of stability predictions."""
        if not self.predictions:
            return {'status': 'No predictions'}
        
        probabilities = [p.probability for p in self.predictions]
        
        return {
            'n_predictions': len(self.predictions),
            'mean_probability': round(np.mean(probabilities), 2),
            'min_probability': round(min(probabilities), 2),
            'max_probability': round(max(probabilities), 2),
            'last_prediction_time': self.predictions[-1].time,
            'last_prediction_prob': self.predictions[-1].probability
        }


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Stability Envelope Prediction (APSEP)")
    print("="*60)
    
    sep = StabilityEnvelopePredictor()
    
    # Current state (from Vikos Gorge example)
    current = (2.8, 3.2, 0.087)
    
    # Predict envelope for 2 years
    envelope = sep.predict_envelope(current, 2.0)
    print(f"\nStability Envelope (2 years):")
    print(f"  V: [{envelope.v_min:.2f}, {envelope.v_max:.2f}] m")
    print(f"  L: [{envelope.l_min:.2f}, {envelope.l_max:.2f}] m")
    print(f"  F: [{envelope.f_min:.3f}, {envelope.f_max:.3f}] mm")
    print(f"  Probability: {envelope.probability:.0%}")
    
    # Predict safe time
    safe = sep.predict_safe_time(current)
    print(f"\nSafe Exploration Time:")
    print(f"  {safe['safe_time_hours']} hours ({safe['safe_time_years']} years)")
    
    print("\n" + "="*60)
