"""
VLF Coupling Dynamics - Analysis of interactions between V, L, and F parameters.
Implements the coupled oscillator model from the research paper.

Author: Samir Baladi
"""

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

logger = logging.getLogger(__name__)

@dataclass
class CouplingCoefficients:
    """Represents the coupling coefficients between V, L, and F."""
    k_vf: float  # V influenced by F
    k_lf: float  # L influenced by F
    k_fv: float  # F influenced by V
    k_fl: float  # F influenced by L


class CouplingDynamics:
    """
    Analyzes coupling dynamics between V, L, and F parameters.
    
    Default coefficients from paper:
    k_vf = 0.23  (Structural stability influenced by flowstone reinforcement)
    k_lf = -0.18 (Lateral dissolution affected by flowstone obstruction)
    k_fv = 0.31  (Deposition rate controlled by passage geometry)
    k_fl = -0.25 (Deposition pattern guided by water flow direction)
    """
    
    def __init__(self,
                 k_vf: float = 0.23,
                 k_lf: float = -0.18,
                 k_fv: float = 0.31,
                 k_fl: float = -0.25):
        """
        Initialize coupling dynamics analyzer.
        
        Args:
            k_vf: V-F coupling coefficient
            k_lf: L-F coupling coefficient
            k_fv: F-V coupling coefficient
            k_fl: F-L coupling coefficient
        """
        self.coefficients = CouplingCoefficients(
            k_vf=k_vf,
            k_lf=k_lf,
            k_fv=k_fv,
            k_fl=k_fl
        )
        
        self.V_history = []
        self.L_history = []
        self.F_history = []
        self.time_history = []
        
        logger.info(f"Initialized CouplingDynamics with coefficients:")
        logger.info(f"  k_vf={k_vf}, k_lf={k_lf}, k_fv={k_fv}, k_fl={k_fl}")

    def add_state(self, 
                  timestamp: float,
                  V: float,
                  L: float,
                  F: float) -> None:
        """
        Add a VLF state at a given time.
        
        Args:
            timestamp: Time of measurement (years)
            V: Vertical passage value
            L: Lateral conduit value
            F: Flowstone value
        """
        self.time_history.append(timestamp)
        self.V_history.append(V)
        self.L_history.append(L)
        self.F_history.append(F)
        logger.debug(f"Added VLF state at t={timestamp}: V={V:.2f}, L={L:.2f}, F={F:.3f}")

    def calculate_coupling_strengths(self) -> Dict[str, float]:
        """
        Calculate current coupling strengths based on historical data.
        
        Returns:
            Dictionary with coupling strengths
        """
        if len(self.V_history) < 10:
            logger.warning("Insufficient data for coupling strength calculation")
            return {
                'V_F': self.coefficients.k_vf,
                'L_F': self.coefficients.k_lf,
                'F_V': self.coefficients.k_fv,
                'F_L': self.coefficients.k_fl
            }
        
        # Convert to numpy arrays
        V = np.array(self.V_history)
        L = np.array(self.L_history)
        F = np.array(self.F_history)
        
        # Calculate cross-correlations
        corr_VF = np.corrcoef(V, F)[0, 1]
        corr_LF = np.corrcoef(L, F)[0, 1]
        
        # Calculate time-lagged correlations for directionality
        lags = range(-5, 6)
        
        # V leading F
        vf_lags = [np.corrcoef(V[:-lag] if lag < 0 else V[lag:], 
                               F[lag:] if lag < 0 else F[:-lag])[0, 1] 
                  for lag in lags if abs(lag) < len(V)//2]
        
        # L leading F
        lf_lags = [np.corrcoef(L[:-lag] if lag < 0 else L[lag:], 
                               F[lag:] if lag < 0 else F[:-lag])[0, 1] 
                  for lag in lags if abs(lag) < len(L)//2]
        
        # Adjust coefficients based on observed correlations
        # This is a simplified adaptation
        k_vf_observed = self.coefficients.k_vf * (1 + 0.5 * corr_VF)
        k_lf_observed = self.coefficients.k_lf * (1 + 0.5 * corr_LF)
        
        return {
            'V_F': round(k_vf_observed, 3),
            'L_F': round(k_lf_observed, 3),
            'F_V': round(self.coefficients.k_fv, 3),
            'F_L': round(self.coefficients.k_fl, 3),
            'correlation_VF': round(corr_VF, 3),
            'correlation_LF': round(corr_LF, 3),
            'n_samples': len(V)
        }

    def compute_derivatives(self, 
                           V: float, 
                           L: float, 
                           F: float,
                           mu_v: float = 0.15,
                           mu_l: float = 0.10,
                           mu_f: float = 0.05) -> Dict[str, float]:
        """
        Compute time derivatives using coupled 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
        
        Args:
            V, L, F: Current state values
            mu_v, mu_l, mu_f: Non-linearity parameters
            
        Returns:
            Dictionary with derivatives
        """
        dVdt = L - mu_v * (V**2 - 1) * V + self.coefficients.k_vf * F
        dLdt = -V + mu_l * (1 - L**2) * L + self.coefficients.k_lf * F
        dFdt = (-self.coefficients.k_fv * V - 
                self.coefficients.k_fl * L + 
                mu_f * (1 - F**2) * F)
        
        return {
            'dVdt': round(dVdt, 4),
            'dLdt': round(dLdt, 4),
            'dFdt': round(dFdt, 4),
            'magnitude': round(np.sqrt(dVdt**2 + dLdt**2 + dFdt**2), 4)
        }

    def analyze_equilibrium(self) -> Dict[str, Union[float, str]]:
        """
        Analyze equilibrium points of the coupled system.
        
        Returns:
            Dictionary with equilibrium analysis
        """
        if len(self.V_history) < 2:
            return {'status': 'Insufficient data'}
        
        # Find fixed points where derivatives are zero
        # This is a simplified grid search
        
        V_range = np.linspace(0, 5, 20)
        L_range = np.linspace(0, 5, 20)
        F_range = np.linspace(0, 1, 10)
        
        min_derivative = float('inf')
        equilibrium = None
        
        for V in V_range:
            for L in L_range:
                for F in F_range:
                    derivs = self.compute_derivatives(V, L, F)
                    mag = derivs['magnitude']
                    
                    if mag < min_derivative:
                        min_derivative = mag
                        equilibrium = (V, L, F)
        
        if equilibrium:
            V_eq, L_eq, F_eq = equilibrium
            
            # Analyze stability of equilibrium
            # Linearize around equilibrium
            stability = self._analyze_linear_stability(V_eq, L_eq, F_eq)
            
            return {
                'V_eq': round(V_eq, 2),
                'L_eq': round(L_eq, 2),
                'F_eq': round(F_eq, 3),
                'derivative_magnitude': round(min_derivative, 4),
                'stability': stability['type'],
                'eigenvalues': stability['eigenvalues']
            }
        
        return {'status': 'No equilibrium found'}

    def _analyze_linear_stability(self, 
                                  V: float, 
                                  L: float, 
                                  F: float,
                                  mu_v: float = 0.15,
                                  mu_l: float = 0.10,
                                  mu_f: float = 0.05) -> Dict[str, Union[str, List]]:
        """
        Analyze linear stability around an equilibrium point.
        """
        # Jacobian matrix of partial derivatives
        # This is a simplified version
        J = np.array([
            [-mu_v * (3*V**2 - 1), 1, self.coefficients.k_vf],
            [-1, mu_l * (1 - 3*L**2), self.coefficients.k_lf],
            [-self.coefficients.k_fv, -self.coefficients.k_fl, -mu_f * (3*F**2 - 1)]
        ])
        
        # Calculate eigenvalues
        eigenvalues = np.linalg.eigvals(J)
        
        # Determine stability type
        real_parts = np.real(eigenvalues)
        if np.all(real_parts < 0):
            stability = 'Stable node'
        elif np.all(real_parts > 0):
            stability = 'Unstable node'
        elif np.any(real_parts > 0) and np.any(real_parts < 0):
            stability = 'Saddle point'
        elif np.any(np.abs(np.imag(eigenvalues)) > 0):
            if np.all(real_parts < 0):
                stability = 'Stable spiral'
            elif np.all(real_parts > 0):
                stability = 'Unstable spiral'
            else:
                stability = 'Center'
        else:
            stability = 'Unknown'
        
        return {
            'type': stability,
            'eigenvalues': [round(ev, 4) for ev in eigenvalues]
        }

    def get_physical_interpretation(self) -> Dict[str, str]:
        """
        Get physical interpretation of coupling coefficients.
        
        Returns:
            Dictionary with interpretations
        """
        return {
            'k_vf': 'Structural stability reinforced by flowstone deposition',
            'k_lf': 'Lateral dissolution obstructed by flowstone growth',
            'k_fv': 'Deposition rate enhanced by vertical passage geometry',
            'k_fl': 'Deposition patterns guided by lateral water flow',
            'k_vf_sign': 'Positive: flowstone stabilizes structure',
            'k_lf_sign': 'Negative: flowstone constricts passages',
            'k_fv_sign': 'Positive: taller passages increase deposition',
            'k_fl_sign': 'Negative: water flow redirects deposition'
        }

    def calculate_optimal_ratio(self) -> Dict[str, float]:
        """
        Calculate optimal V:L:F ratio for maximum stability.
        
        Returns:
            Dictionary with optimal ratios
        """
        # From paper: optimal coupling occurs when system is balanced
        
        # Search for configuration minimizing derivative magnitude
        V_opt, L_opt, F_opt = 2.5, 3.0, 0.08  # From paper
        
        return {
            'V_optimal': V_opt,
            'L_optimal': L_opt,
            'F_optimal': F_opt,
            'V_L_ratio': round(V_opt / L_opt, 3),
            'V_F_ratio': round(V_opt / F_opt, 1),
            'L_F_ratio': round(L_opt / F_opt, 1),
            'interpretation': 'V:L:F = 2.5:3.0:0.08 for optimal stability'
        }

    def detect_coupling_regime(self) -> Dict[str, str]:
        """
        Detect current coupling regime based on historical data.
        
        Returns:
            Dictionary with regime classification
        """
        if len(self.V_history) < 10:
            return {'regime': 'Insufficient data'}
        
        strengths = self.calculate_coupling_strengths()
        
        # Classify based on coupling strengths
        if strengths['V_F'] > 0.3 and abs(strengths['L_F']) < 0.1:
            regime = "V-dominated"
            description = "Vertical passage controls dynamics"
        elif abs(strengths['L_F']) > 0.2 and strengths['V_F'] < 0.2:
            regime = "L-dominated"
            description = "Lateral hydrology controls dynamics"
        elif strengths['F_V'] > 0.3 and strengths['F_L'] > 0.2:
            regime = "F-dominated"
            description = "Flowstone deposition controls dynamics"
        elif (strengths['V_F'] > 0.2 and abs(strengths['L_F']) > 0.15):
            regime = "Balanced V-L-F"
            description = "All three dimensions actively coupled"
        else:
            regime = "Weak coupling"
            description = "Dimensions weakly coupled"
        
        return {
            'regime': regime,
            'description': description,
            'V_coupling': strengths['V_F'],
            'L_coupling': strengths['L_F'],
            'F_coupling': strengths['F_V']
        }

    def summary(self) -> Dict[str, Union[float, str, Dict]]:
        """Get summary of coupling dynamics."""
        return {
            'coefficients': {
                'k_vf': self.coefficients.k_vf,
                'k_lf': self.coefficients.k_lf,
                'k_fv': self.coefficients.k_fv,
                'k_fl': self.coefficients.k_fl
            },
            'interpretations': self.get_physical_interpretation(),
            'n_observations': len(self.V_history),
            'time_span': (self.time_history[-1] - self.time_history[0]) if self.time_history else 0,
            'current_regime': self.detect_coupling_regime() if self.V_history else {'regime': 'No data'},
            'optimal_ratio': self.calculate_optimal_ratio()
        }


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Coupling Dynamics Analysis")
    print("="*60)
    
    cd = CouplingDynamics()
    
    # Add some historical data
    times = np.linspace(2015, 2025, 20)
    for i, t in enumerate(times):
        V = 2.5 + 0.2 * np.sin(2 * np.pi * 0.05 * i)
        L = 3.0 + 0.3 * np.cos(2 * np.pi * 0.05 * i)
        F = 0.08 + 0.01 * np.sin(2 * np.pi * 0.1 * i)
        cd.add_state(t, V, L, F)
    
    # Calculate coupling strengths
    strengths = cd.calculate_coupling_strengths()
    print(f"\nCoupling Strengths:")
    print(f"  - V-F: {strengths['V_F']}")
    print(f"  - L-F: {strengths['L_F']}")
    print(f"  - F-V: {strengths['F_V']}")
    print(f"  - F-L: {strengths['F_L']}")
    
    # Compute derivatives for a sample state
    derivs = cd.compute_derivatives(V=2.8, L=3.2, F=0.087)
    print(f"\nDerivatives:")
    print(f"  - dV/dt: {derivs['dVdt']}")
    print(f"  - dL/dt: {derivs['dLdt']}")
    print(f"  - dF/dt: {derivs['dFdt']}")
    
    # Analyze equilibrium
    equilibrium = cd.analyze_equilibrium()
    if 'V_eq' in equilibrium:
        print(f"\nEquilibrium Point:")
        print(f"  - V_eq: {equilibrium['V_eq']} m")
        print(f"  - L_eq: {equilibrium['L_eq']} m")
        print(f"  - F_eq: {equilibrium['F_eq']} mm")
        print(f"  - Stability: {equilibrium['stability']}")
    
    # Detect regime
    regime = cd.detect_coupling_regime()
    print(f"\nCoupling Regime:")
    print(f"  - Regime: {regime['regime']}")
    print(f"  - Description: {regime['description']}")
    
    # Physical interpretations
    interpretations = cd.get_physical_interpretation()
    print(f"\nPhysical Interpretations:")
    for key, value in interpretations.items():
        print(f"  - {key}: {value}")
    
    print("\n" + "="*60)
