"""
Phase space reconstruction and analysis for cave passage dynamics.
Provides tools for attractor reconstruction and dimensionality analysis.

Author: Samir Baladi
"""

import numpy as np
from typing import Tuple, Optional, Dict, List
from scipy.spatial import KDTree
from scipy import signal
import logging

logger = logging.getLogger(__name__)

class PhaseSpaceReconstructor:
    """
    Reconstructs and analyzes phase space from time series data using
    time-delay embedding techniques.
    """
    
    def __init__(self, 
                 max_dimension: int = 5,
                 time_delay: Optional[int] = None,
                 method: str = 'mutual_information'):
        """
        Initialize the phase space reconstructor.
        
        Args:
            max_dimension: Maximum embedding dimension to consider
            time_delay: Time delay for embedding (auto-calculated if None)
            method: Method for determining time delay ('mutual_information' or 'autocorrelation')
        """
        self.max_dimension = max_dimension
        self.time_delay = time_delay
        self.method = method
        logger.info(f"Initialized PhaseSpaceReconstructor with max_dim={max_dimension}")

    def reconstruct(self, 
                   data: np.ndarray,
                   dt: float = 1.0) -> Dict[str, np.ndarray]:
        """
        Reconstruct phase space from time series data.
        
        Args:
            data: Input time series
            dt: Time step between samples (days)
            
        Returns:
            Dictionary with reconstructed phase space and metadata
        """
        logger.info(f"Reconstructing phase space from time series of length {len(data)}")
        
        # Determine optimal time delay if not provided
        if self.time_delay is None:
            self.time_delay = self._estimate_time_delay(data, dt)
            logger.info(f"Estimated time delay: {self.time_delay} samples ({self.time_delay*dt:.1f} days)")
        
        # Determine optimal embedding dimension
        embedding_dim = self._estimate_embedding_dimension(data)
        logger.info(f"Estimated embedding dimension: {embedding_dim}")
        
        # Reconstruct phase space
        phase_space = self._time_delay_embedding(data, embedding_dim, self.time_delay)
        
        return {
            'phase_space': phase_space,
            'embedding_dimension': embedding_dim,
            'time_delay': self.time_delay,
            'time_delay_days': self.time_delay * dt,
            'n_points': len(phase_space)
        }

    def _estimate_time_delay(self, data: np.ndarray, dt: float) -> int:
        """
        Estimate optimal time delay using mutual information or autocorrelation.
        """
        if self.method == 'mutual_information':
            return self._mutual_information_delay(data)
        else:
            return self._autocorrelation_delay(data)

    def _mutual_information_delay(self, data: np.ndarray, max_delay: int = 50) -> int:
        """
        Estimate time delay using first minimum of mutual information.
        """
        # Simplified: use first minimum of autocorrelation as proxy
        # In practice, use proper mutual information calculation
        return self._autocorrelation_delay(data)

    def _autocorrelation_delay(self, data: np.ndarray, max_delay: int = 50) -> int:
        """
        Estimate time delay using autocorrelation function.
        """
        # Normalize data
        data_norm = (data - np.mean(data)) / np.std(data)
        
        # Calculate autocorrelation for different lags
        autocorr = np.correlate(data_norm, data_norm, mode='full')
        autocorr = autocorr[len(autocorr)//2:]
        
        # Find first lag where autocorrelation drops below 1/e
        threshold = 1.0 / np.e
        for delay in range(1, min(max_delay, len(autocorr))):
            if abs(autocorr[delay]) < threshold:
                return delay
        
        return max_delay // 4  # Default if not found

    def _estimate_embedding_dimension(self, data: np.ndarray) -> int:
        """
        Estimate optimal embedding dimension using false nearest neighbors.
        """
        # Simplified: use correlation dimension approach
        # In practice, use proper false nearest neighbors algorithm
        
        # Try different dimensions and look for saturation
        dimensions = range(1, self.max_dimension + 1)
        corr_dims = []
        
        for dim in dimensions:
            if self.time_delay is None:
                continue
            phase = self._time_delay_embedding(data, dim, self.time_delay)
            if len(phase) < 10:
                corr_dims.append(0.0)
                continue
            
            # Approximate correlation dimension
            corr_dim = self._approximate_correlation_dimension(phase)
            corr_dims.append(corr_dim)
        
        # Find where correlation dimension saturates
        if len(corr_dims) < 2:
            return 3  # Default from paper
        
        # Look for plateau
        for i in range(1, len(corr_dims)):
            if abs(corr_dims[i] - corr_dims[i-1]) < 0.1:
                return i
        
        return 3  # Default embedding dimension from paper

    def _time_delay_embedding(self, 
                             data: np.ndarray, 
                             dim: int, 
                             delay: int) -> np.ndarray:
        """
        Perform time-delay embedding of the data.
        """
        n_points = len(data) - (dim - 1) * delay
        if n_points <= 0:
            return np.array([])
        
        phase_space = []
        for i in range(n_points):
            point = [data[i + j * delay] for j in range(dim)]
            phase_space.append(point)
        
        return np.array(phase_space)

    def _approximate_correlation_dimension(self, phase_space: np.ndarray) -> float:
        """
        Approximate correlation dimension of the phase space.
        """
        if len(phase_space) < 10:
            return 0.0
        
        # Use small radius approximation
        tree = KDTree(phase_space)
        
        # Count neighbors within small radius
        r = 0.1 * np.std(phase_space)  # Small radius
        neighbors = []
        
        for i in range(len(phase_space)):
            indices = tree.query_ball_point(phase_space[i], r)
            neighbors.append(len(indices) - 1)  # Exclude self
        
        avg_neighbors = np.mean(neighbors)
        
        if avg_neighbors == 0:
            return 0.0
        
        # Correlation dimension ≈ log(avg_neighbors) / log(r)
        # Simplified scaling
        return np.log(avg_neighbors + 1) / np.log(r + 1)

    def calculate_lyapunov_exponents(self, 
                                    phase_space: np.ndarray, 
                                    dt: float) -> Dict[str, float]:
        """
        Calculate Lyapunov exponents from phase space trajectory.
        
        Args:
            phase_space: Reconstructed phase space points
            dt: Time step between points (days)
            
        Returns:
            Dictionary with Lyapunov exponents
        """
        if len(phase_space) < 10:
            return {'lambda_max': 0.0, 'spectrum': []}
        
        # Simplified: use maximum Lyapunov exponent approximation
        # In practice, use proper algorithm (Wolf et al., 1985)
        
        tree = KDTree(phase_space)
        divergence_rates = []
        
        # Track divergence of nearby trajectories
        for i in range(0, len(phase_space) - 10, 5):
            # Find nearest neighbor (excluding temporally close points)
            distances, indices = tree.query(phase_space[i], k=5)
            
            for j, idx in enumerate(indices[1:], 1):
                if abs(idx - i) > 10:  # Not temporally correlated
                    # Track divergence over next few steps
                    d0 = distances[j]
                    if d0 < 1e-6:
                        continue
                    
                    max_steps = min(10, len(phase_space) - max(i, idx))
                    if max_steps < 3:
                        continue
                    
                    for step in range(1, max_steps):
                        if i + step >= len(phase_space) or idx + step >= len(phase_space):
                            break
                        
                        d = np.linalg.norm(phase_space[i + step] - phase_space[idx + step])
                        if d > 0:
                            divergence_rates.append(np.log(d / d0) / (step * dt))
        
        if not divergence_rates:
            return {'lambda_max': 0.0, 'spectrum': []}
        
        lambda_max = np.mean(divergence_rates)
        
        # For limit cycle, spectrum should have one zero, others negative
        # Simplified spectrum
        spectrum = [lambda_max, 0.0, -0.5 * lambda_max]
        
        return {
            'lambda_max': float(lambda_max),
            'spectrum': [float(l) for l in spectrum],
            'interpretation': self._interpret_lyapunov(spectrum)
        }

    def _interpret_lyapunov(self, spectrum: List[float]) -> str:
        """
        Interpret the Lyapunov spectrum.
        """
        positive = sum(1 for l in spectrum if l > 0.01)
        zero = sum(1 for l in spectrum if abs(l) <= 0.01)
        negative = sum(1 for l in spectrum if l < -0.01)
        
        if positive > 0:
            return "Chaotic attractor"
        elif zero == 1 and negative >= 0:
            return "Limit cycle (periodic)"
        elif zero == 0 and negative > 0:
            return "Fixed point"
        else:
            return "Unknown dynamics"

    def plot_phase_portrait(self, 
                          phase_space: np.ndarray,
                          dimensions: Tuple[int, int] = (0, 1),
                          ax=None,
                          **kwargs):
        """
        Create a 2D phase portrait projection.
        
        Args:
            phase_space: Phase space points
            dimensions: Which dimensions to plot (default: first two)
            ax: Matplotlib axis (creates new if None)
            **kwargs: Additional plot arguments
        
        Returns:
            Matplotlib axis
        """
        try:
            import matplotlib.pyplot as plt
            
            if ax is None:
                _, ax = plt.subplots(figsize=(8, 6))
            
            ax.plot(phase_space[:, dimensions[0]], 
                   phase_space[:, dimensions[1]], 
                   'b-', alpha=0.7, linewidth=0.5, **kwargs)
            ax.plot(phase_space[0, dimensions[0]], 
                   phase_space[0, dimensions[1]], 
                   'go', markersize=8, label='Start')
            ax.plot(phase_space[-1, dimensions[0]], 
                   phase_space[-1, dimensions[1]], 
                   'ro', markersize=8, label='End')
            
            ax.set_xlabel(f'Dimension {dimensions[0] + 1}')
            ax.set_ylabel(f'Dimension {dimensions[1] + 1}')
            ax.set_title('Phase Space Portrait')
            ax.legend()
            ax.grid(True, alpha=0.3)
            
            return ax
            
        except ImportError:
            logger.warning("matplotlib not available for plotting")
            return None


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Phase Space Reconstruction")
    print("="*60)
    
    # Generate sample data (limit cycle + noise)
    t = np.linspace(0, 50, 1000)
    dt = t[1] - t[0]
    V = 2.0 + 0.5 * np.sin(2 * np.pi * 0.05 * t) + 0.1 * np.random.randn(len(t))
    
    # Reconstruct phase space
    reconstructor = PhaseSpaceReconstructor()
    result = reconstructor.reconstruct(V, dt)
    
    print(f"\nReconstruction Results:")
    print(f"  - Embedding Dimension: {result['embedding_dimension']}")
    print(f"  - Time Delay: {result['time_delay']} samples ({result['time_delay_days']:.1f} days)")
    print(f"  - Phase Space Points: {result['n_points']}")
    
    # Calculate Lyapunov exponents
    lyapunov = reconstructor.calculate_lyapunov_exponents(result['phase_space'], dt)
    print(f"\nLyapunov Analysis:")
    print(f"  - Maximum Exponent: {lyapunov['lambda_max']:.4f}")
    print(f"  - Spectrum: {[f'{l:.4f}' for l in lyapunov['spectrum']]}")
    print(f"  - Interpretation: {lyapunov['interpretation']}")
    
    print("\n" + "="*60)
