"""
Van der Pol Oscillator models for karst passage dynamics.
Based on the research paper: "Limit Cycle Hydrological Dynamics as a Framework 
for Adaptive Cave Passage Safety Protocols"

Equation: d²V/dt² - μ(1 - V²) dV/dt + ω₀²V = F_ext(t)

Author: Samir Baladi
"""

import numpy as np
from scipy.integrate import solve_ivp
from typing import Tuple, Optional, Dict, Union, Callable
import logging

logger = logging.getLogger(__name__)

class VanDerPolOscillator:
    """
    Implements the Van der Pol oscillator model for a single dimension (e.g., Vertical Passage V).
    
    Attributes:
        mu (float): Non-linearity parameter. Default 0.15 for Mediterranean karst.
        omega0 (float): Natural frequency. Default 0.05 rad/yr.
    """
    
    def __init__(self, mu: float = 0.15, omega0: float = 0.05):
        """
        Initialize the Van der Pol oscillator with calibrated parameters.
        
        Args:
            mu: Non-linearity parameter from paper (0.15 for Mediterranean karst)
            omega0: Natural frequency (0.05 rad/yr for typical systems)
        """
        self.mu = mu
        self.omega0 = omega0
        self.omega0_sq = omega0**2
        logger.info(f"Initialized VanDerPolOscillator with μ={mu}, ω₀={omega0}")

    def system_equations(self, t: float, y: np.ndarray, F_ext: float = 0.0) -> np.ndarray:
        """
        System of first-order ODEs for the Van der Pol oscillator.
        
        Args:
            t: Time
            y: State vector [V, dV/dt]
            F_ext: External forcing value at time t
            
        Returns:
            Derivative vector [dV/dt, d²V/dt²]
        """
        V, V_dot = y
        V_ddot = self.mu * (1 - V**2) * V_dot - self.omega0_sq * V + F_ext
        return np.array([V_dot, V_ddot])

    def simulate(self, 
                 t_span: Tuple[float, float], 
                 y0: Tuple[float, float], 
                 t_eval: Optional[np.ndarray] = None,
                 F_ext_func: Optional[Callable] = None) -> Dict[str, np.ndarray]:
        """
        Simulate the oscillator over a given time span.
        
        Args:
            t_span: Tuple of (start_time, end_time) in years
            y0: Initial conditions (V0, dVdt0)
            t_eval: Optional array of time points for evaluation
            F_ext_func: Optional function F_ext(t) returning external forcing
            
        Returns:
            Dictionary with keys 't', 'V', 'V_dot'
        """
        if F_ext_func is None:
            F_ext_func = lambda t: 0.0
        
        def rhs(t, y):
            return self.system_equations(t, y, F_ext_func(t))
        
        logger.info(f"Simulating Van der Pol oscillator from t={t_span[0]} to {t_span[1]} years")
        sol = solve_ivp(rhs, t_span, y0, t_eval=t_eval, method='RK45', rtol=1e-6, atol=1e-9)
        
        if not sol.success:
            logger.error(f"Simulation failed: {sol.message}")
            raise RuntimeError(f"Simulation failed: {sol.message}")
        
        logger.info(f"Simulation successful: {len(sol.t)} time steps")
        
        return {
            't': sol.t,
            'V': sol.y[0, :],
            'V_dot': sol.y[1, :]
        }

    def get_limit_cycle_amplitude(self) -> float:
        """
        Calculate the approximate limit cycle amplitude.
        
        From paper: A_lc = 2 * sqrt(mu / omega0^2)
        
        Returns:
            Limit cycle amplitude (same units as V)
        """
        amplitude = 2 * np.sqrt(self.mu / self.omega0_sq)
        logger.info(f"Limit cycle amplitude: {amplitude:.3f}")
        return amplitude

    def get_limit_cycle_period(self) -> float:
        """
        Estimate the period of the limit cycle.
        
        Returns:
            Approximate period in years
        """
        period = 2 * np.pi / self.omega0
        logger.info(f"Limit cycle period: {period:.1f} years")
        return period


class CoupledVLFOscillator:
    """
    Implements the coupled three-dimensional model for V, L, F.
    
    Equations:
    dV/dt = L - μ_v(V² - 1)V + k_vf F
    dL/dt = -V + μ_l(1 - L²)L + k_lf F
    dF/dt = -k_fv V - k_fl L + μ_f(1 - F²)F
    
    Where:
    V: Vertical passage control (structural)
    L: Lateral conduit control (hydrological)
    F: Flowstone deposition (sedimentary)
    """
    
    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 the coupled oscillator with calibrated coefficients from the paper.
        
        Args:
            mu_v, mu_l, mu_f: Non-linearity parameters for V, L, F
            k_vf, k_lf, k_fv, k_fl: Coupling coefficients
        """
        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
        
        logger.info(f"Initialized CoupledVLFOscillator with coefficients: "
                   f"k_vf={k_vf}, k_lf={k_lf}, k_fv={k_fv}, k_fl={k_fl}")

    def system_equations(self, t: float, y: np.ndarray) -> np.ndarray:
        """
        System of first-order ODEs for the coupled VLF model.
        
        Args:
            t: Time
            y: State vector [V, L, F]
        
        Returns:
            Derivative vector [dV/dt, dL/dt, dF/dt]
        """
        V, L, F = y
        
        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 np.array([dVdt, dLdt, dFdt])

    def simulate(self,
                 t_span: Tuple[float, float],
                 y0: Tuple[float, float, float],
                 t_eval: Optional[np.ndarray] = None) -> Dict[str, np.ndarray]:
        """
        Simulate the coupled system over a given time span.
        
        Args:
            t_span: Tuple of (start_time, end_time) in years
            y0: Initial conditions (V0, L0, F0)
            t_eval: Optional array of time points for evaluation
        
        Returns:
            Dictionary with keys 't', 'V', 'L', 'F'
        """
        logger.info(f"Simulating coupled VLF system from t={t_span[0]} to {t_span[1]} years")
        logger.info(f"Initial conditions: V={y0[0]:.2f}, L={y0[1]:.2f}, F={y0[2]:.3f}")
        
        sol = solve_ivp(self.system_equations, t_span, y0, t_eval=t_eval, 
                       method='RK45', rtol=1e-6, atol=1e-9)
        
        if not sol.success:
            logger.error(f"Simulation failed: {sol.message}")
            raise RuntimeError(f"Simulation failed: {sol.message}")
        
        logger.info(f"Simulation successful: {len(sol.t)} time steps")
        
        return {
            't': sol.t,
            'V': sol.y[0, :],
            'L': sol.y[1, :],
            'F': sol.y[2, :]
        }

    def get_stability_envelope(self, state: Tuple[float, float, float]) -> Dict[str, Union[float, str]]:
        """
        Calculate the stability envelope for a given state based on proximity to equilibrium.
        
        Args:
            state: Tuple (V, L, F) representing the current passage state
        
        Returns:
            Dictionary with stability score and level
        """
        V, L, F = state
        
        # Check for extreme values indicating immediate danger
        if abs(V) > 10 or abs(L) > 10 or abs(F) > 10:
            return {
                "score": 10.0,
                "level": "CRITICAL",
                "alert": "EMERGENCY",
                "color": "⚫",
                "action": "Immediate Evacuation"
            }
        
        # Calculate derivatives to assess change potential
        dVdt = abs(L - self.mu_v * (V**2 - 1) * V + self.k_vf * F)
        dLdt = abs(-V + self.mu_l * (1 - L**2) * L + self.k_lf * F)
        dFdt = abs(-self.k_fv * V - self.k_fl * L + self.mu_f * (1 - F**2) * F)
        
        # Combine derivatives into a "change potential" score (0-100, higher is more unstable)
        change_potential = (dVdt + dLdt + dFdt) * 10
        score = max(0.0, min(100.0, 100.0 - change_potential))
        
        # Classify stability level based on paper thresholds
        if score >= 80:
            level = "STABLE"
            alert = "None"
            color = "🟢"
            action = "Standard Survey Mode"
            max_hours = "Unrestricted"
        elif score >= 60:
            level = "CAUTION"
            alert = "Advisory"
            color = "🟡"
            action = "Enhanced Monitoring"
            max_hours = "8 hours/shift"
        elif score >= 40:
            level = "UNSTABLE"
            alert = "Warning"
            color = "🟠"
            action = "CDZ Protocol Active"
            max_hours = "6 hours/shift"
        elif score >= 20:
            level = "HIGH RISK"
            alert = "Alert"
            color = "🔴"
            action = "Evacuation Preparation"
            max_hours = "4 hours/shift"
        else:
            level = "CRITICAL"
            alert = "Emergency"
            color = "⚫"
            action = "Immediate Evacuation"
            max_hours = "0 hours (exit)"
        
        return {
            "score": round(score, 1),
            "level": level,
            "alert": alert,
            "color": color,
            "action": action,
            "max_hours": max_hours
        }

    def get_coupling_interpretation(self) -> Dict[str, str]:
        """
        Get the physical interpretation of coupling coefficients.
        
        Returns:
            Dictionary with coupling interpretations
        """
        return {
            "k_vf": "Structural stability influenced by flowstone reinforcement",
            "k_lf": "Lateral dissolution affected by flowstone obstruction",
            "k_fv": "Deposition rate controlled by passage geometry",
            "k_fl": "Deposition pattern guided by water flow direction"
        }


# Example usage
if __name__ == "__main__":
    # Setup logging
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Van der Pol Oscillator Models")
    print("="*60)
    
    # Example 1: Single Van der Pol oscillator
    print("\n1. Simulating Single Van der Pol Oscillator (Vertical Passage V):")
    vdp = VanDerPolOscillator()
    
    t_span = (0, 200)
    y0 = (0.5, 0.0)
    t_eval = np.linspace(0, 200, 1000)
    
    result_vdp = vdp.simulate(t_span, y0, t_eval=t_eval)
    
    print(f"   - Limit Cycle Amplitude: {vdp.get_limit_cycle_amplitude():.3f}")
    print(f"   - Limit Cycle Period: {vdp.get_limit_cycle_period():.1f} years")
    print(f"   - Final V value: {result_vdp['V'][-1]:.3f}")
    
    # Example 2: Coupled VLF model
    print("\n2. Simulating Coupled VLF Model:")
    coupled = CoupledVLFOscillator()
    
    y0_vlf = (2.5, 3.0, 0.08)  # Typical stable passage
    t_span_vlf = (0, 50)
    t_eval_vlf = np.linspace(0, 50, 500)
    
    result_vlf = coupled.simulate(t_span_vlf, y0_vlf, t_eval=t_eval_vlf)
    
    print(f"   - Final state after 50 yrs: V={result_vlf['V'][-1]:.2f}, "
          f"L={result_vlf['L'][-1]:.2f}, F={result_vlf['F'][-1]:.3f}")
    
    # Example 3: Stability assessment
    print("\n3. Stability Envelope Assessment:")
    test_state = (2.8, 3.2, 0.087)  # Vikos Gorge example
    stability = coupled.get_stability_envelope(test_state)
    
    print(f"   - State (V, L, F): {test_state}")
    print(f"   - Stability Score: {stability['score']}/100")
    print(f"   - Level: {stability['color']} {stability['level']}")
    print(f"   - Action: {stability['action']}")
    print(f"   - Max Hours: {stability['max_hours']}")
    
    print("\n" + "="*60)
    print("Models loaded successfully")
    print("="*60)
