"""
Phase portrait visualization for dynamical systems analysis.
Creates phase space plots, Poincaré sections, and attractor visualizations.

Author: Samir Baladi
"""

import numpy as np
from typing import Dict, List, Tuple, Optional, Union
import logging

logger = logging.getLogger(__name__)

try:
    import matplotlib.pyplot as plt
    from matplotlib import cm
    MATPLOTLIB_AVAILABLE = True
except ImportError:
    MATPLOTLIB_AVAILABLE = False
    logger.warning("Matplotlib not available - phase portrait functionality limited")

try:
    import plotly.graph_objects as go
    PLOTLY_AVAILABLE = True
except ImportError:
    PLOTLY_AVAILABLE = False


class PhasePortraits:
    """
    Creates phase portraits for dynamical systems analysis.
    """
    
    def __init__(self):
        """Initialize phase portrait visualizer."""
        self.figures = []
        logger.info("Initialized PhasePortraits")

    def create_2d_portrait(self,
                          x: np.ndarray,
                          y: np.ndarray,
                          xlabel: str = "X",
                          ylabel: str = "Y",
                          title: str = "Phase Portrait",
                          color_by_time: bool = True,
                          show_start_end: bool = True,
                          ax=None) -> Optional[plt.Figure]:
        """
        Create 2D phase portrait.
        
        Args:
            x, y: Trajectory coordinates
            xlabel, ylabel: Axis labels
            title: Plot title
            color_by_time: Color trajectory by time
            show_start_end: Show start and end points
            ax: Matplotlib axis (creates new if None)
            
        Returns:
            Matplotlib figure or None if matplotlib unavailable
        """
        if not MATPLOTLIB_AVAILABLE:
            return None
        
        if ax is None:
            fig, ax = plt.subplots(figsize=(10, 8))
        else:
            fig = ax.figure
        
        if color_by_time:
            # Color by time
            colors = np.arange(len(x))
            scatter = ax.scatter(x, y, c=colors, cmap='viridis', 
                               s=10, alpha=0.6, label='Trajectory')
            plt.colorbar(scatter, ax=ax, label='Time')
        else:
            ax.plot(x, y, 'b-', alpha=0.7, linewidth=1, label='Trajectory')
        
        if show_start_end:
            ax.plot(x[0], y[0], 'go', markersize=10, label='Start', zorder=5)
            ax.plot(x[-1], y[-1], 'ro', markersize=10, label='Current', zorder=5)
        
        ax.set_xlabel(xlabel)
        ax.set_ylabel(ylabel)
        ax.set_title(title)
        ax.grid(True, alpha=0.3)
        ax.legend()
        
        self.figures.append(fig)
        return fig

    def create_3d_portrait(self,
                          x: np.ndarray,
                          y: np.ndarray,
                          z: np.ndarray,
                          xlabel: str = "X",
                          ylabel: str = "Y",
                          zlabel: str = "Z",
                          title: str = "3D Phase Portrait") -> Optional[go.Figure]:
        """
        Create 3D phase portrait using Plotly.
        
        Args:
            x, y, z: Trajectory coordinates
            xlabel, ylabel, zlabel: Axis labels
            title: Plot title
            
        Returns:
            Plotly figure or None if plotly unavailable
        """
        if not PLOTLY_AVAILABLE:
            return None
        
        # Color by time
        colors = np.arange(len(x))
        
        fig = go.Figure(data=[
            go.Scatter3d(
                x=x, y=y, z=z,
                mode='lines+markers',
                marker=dict(
                    size=2,
                    color=colors,
                    colorscale='Viridis',
                    showscale=True,
                    colorbar=dict(title="Time")
                ),
                line=dict(color='darkblue', width=1)
            )
        ])
        
        # Add start and end points
        fig.add_trace(go.Scatter3d(
            x=[x[0]], y=[y[0]], z=[z[0]],
            mode='markers',
            marker=dict(size=8, color='green'),
            name='Start'
        ))
        
        fig.add_trace(go.Scatter3d(
            x=[x[-1]], y=[y[-1]], z=[z[-1]],
            mode='markers',
            marker=dict(size=8, color='red'),
            name='Current'
        ))
        
        fig.update_layout(
            title=title,
            scene=dict(
                xaxis_title=xlabel,
                yaxis_title=ylabel,
                zaxis_title=zlabel
            ),
            height=600
        )
        
        return fig

    def create_poincare_section(self,
                               x: np.ndarray,
                               y: np.ndarray,
                               threshold: float,
                               plane: str = 'x',
                               title: str = "Poincaré Section") -> Optional[plt.Figure]:
        """
        Create Poincaré section plot.
        
        Args:
            x, y: Trajectory coordinates
            threshold: Section threshold
            plane: Section plane ('x' or 'y')
            title: Plot title
            
        Returns:
            Matplotlib figure or None if matplotlib unavailable
        """
        if not MATPLOTLIB_AVAILABLE:
            return None
        
        # Find crossings
        crossings_x = []
        crossings_y = []
        
        if plane == 'x':
            for i in range(1, len(x)):
                if (x[i-1] - threshold) * (x[i] - threshold) < 0:
                    # Linear interpolation
                    t = (threshold - x[i-1]) / (x[i] - x[i-1])
                    cross_x = threshold
                    cross_y = y[i-1] + t * (y[i] - y[i-1])
                    crossings_x.append(cross_x)
                    crossings_y.append(cross_y)
        else:  # plane == 'y'
            for i in range(1, len(y)):
                if (y[i-1] - threshold) * (y[i] - threshold) < 0:
                    t = (threshold - y[i-1]) / (y[i] - y[i-1])
                    cross_x = x[i-1] + t * (x[i] - x[i-1])
                    cross_y = threshold
                    crossings_x.append(cross_x)
                    crossings_y.append(cross_y)
        
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))
        
        # Phase portrait with section line
        ax1.plot(x, y, 'b-', alpha=0.5, linewidth=0.5)
        if plane == 'x':
            ax1.axvline(x=threshold, color='r', linestyle='--', alpha=0.7)
        else:
            ax1.axhline(y=threshold, color='r', linestyle='--', alpha=0.7)
        ax1.plot(crossings_x, crossings_y, 'ro', markersize=4)
        ax1.set_xlabel('X')
        ax1.set_ylabel('Y')
        ax1.set_title('Trajectory with Section')
        ax1.grid(True, alpha=0.3)
        
        # Poincaré section points
        if crossings_x:
            ax2.scatter(crossings_x, crossings_y, c='red', s=30, alpha=0.6)
            ax2.set_xlabel('X')
            ax2.set_ylabel('Y')
            ax2.set_title(f'Poincaré Section ({len(crossings_x)} crossings)')
            ax2.grid(True, alpha=0.3)
            
            # Add return map if enough points
            if len(crossings_x) > 2:
                ax3 = fig.add_subplot(133)
                ax3.scatter(crossings_x[:-1], crossings_x[1:], c='blue', alpha=0.6)
                ax3.plot([min(crossings_x), max(crossings_x)], 
                        [min(crossings_x), max(crossings_x)], 'k--', alpha=0.5)
                ax3.set_xlabel('X_n')
                ax3.set_ylabel('X_{n+1}')
                ax3.set_title('Return Map')
                ax3.grid(True, alpha=0.3)
        
        plt.suptitle(title)
        plt.tight_layout()
        
        self.figures.append(fig)
        return fig

    def create_attractor_plot(self,
                             x: np.ndarray,
                             y: np.ndarray,
                             z: np.ndarray,
                             title: str = "Attractor Visualization") -> Optional[plt.Figure]:
        """
        Create attractor visualization with density coloring.
        
        Args:
            x, y, z: Trajectory coordinates
            title: Plot title
            
        Returns:
            Matplotlib figure or None if matplotlib unavailable
        """
        if not MATPLOTLIB_AVAILABLE:
            return None
        
        # Calculate point density for coloring
        from scipy.stats import gaussian_kde
        
        points = np.vstack([x, y])
        kde = gaussian_kde(points)
        density = kde.evaluate(points)
        
        fig = plt.figure(figsize=(12, 5))
        
        # 2D projection with density
        ax1 = fig.add_subplot(121)
        scatter = ax1.scatter(x, y, c=density, s=1, cmap='hot', alpha=0.6)
        plt.colorbar(scatter, ax=ax1, label='Density')
        ax1.set_xlabel('X')
        ax1.set_ylabel('Y')
        ax1.set_title('Attractor Density')
        ax1.grid(True, alpha=0.3)
        
        # 3D view if z provided
        if z is not None and len(z) > 0:
            ax2 = fig.add_subplot(122, projection='3d')
            ax2.scatter(x, y, z, c=density, s=1, cmap='hot', alpha=0.6)
            ax2.set_xlabel('X')
            ax2.set_ylabel('Y')
            ax2.set_zlabel('Z')
            ax2.set_title('3D Attractor')
        
        plt.suptitle(title)
        plt.tight_layout()
        
        self.figures.append(fig)
        return fig

    def create_phase_space_reconstruction(self,
                                         data: np.ndarray,
                                         delay: int = 1,
                                         dimension: int = 3,
                                         title: str = "Phase Space Reconstruction") -> Optional[plt.Figure]:
        """
        Create phase space reconstruction using time-delay embedding.
        
        Args:
            data: Time series data
            delay: Time delay
            dimension: Embedding dimension
            title: Plot title
            
        Returns:
            Matplotlib figure or None if matplotlib unavailable
        """
        if not MATPLOTLIB_AVAILABLE:
            return None
        
        # Reconstruct phase space
        n_points = len(data) - (dimension - 1) * delay
        if n_points <= 0:
            return None
        
        phase = []
        for i in range(n_points):
            point = [data[i + j * delay] for j in range(dimension)]
            phase.append(point)
        
        phase = np.array(phase)
        
        if dimension == 2:
            fig, ax = plt.subplots(figsize=(10, 8))
            ax.plot(phase[:, 0], phase[:, 1], 'b-', alpha=0.7, linewidth=0.5)
            ax.set_xlabel(f'X(t)')
            ax.set_ylabel(f'X(t+{delay})')
            ax.set_title(title)
            ax.grid(True, alpha=0.3)
            
        elif dimension >= 3:
            fig = plt.figure(figsize=(12, 5))
            
            ax1 = fig.add_subplot(121)
            ax1.plot(phase[:, 0], phase[:, 1], 'b-', alpha=0.7, linewidth=0.5)
            ax1.set_xlabel(f'X(t)')
            ax1.set_ylabel(f'X(t+{delay})')
            ax1.set_title('2D Projection')
            ax1.grid(True, alpha=0.3)
            
            ax2 = fig.add_subplot(122, projection='3d')
            ax2.plot(phase[:, 0], phase[:, 1], phase[:, 2], 
                    'b-', alpha=0.7, linewidth=0.5)
            ax2.set_xlabel(f'X(t)')
            ax2.set_ylabel(f'X(t+{delay})')
            ax2.set_zlabel(f'X(t+{2*delay})')
            ax2.set_title('3D Reconstruction')
        
        self.figures.append(fig)
        return fig

    def save_figure(self, filename: str, fig=None) -> None:
        """
        Save figure to file.
        
        Args:
            filename: Output filename
            fig: Figure to save (saves last if None)
        """
        if not MATPLOTLIB_AVAILABLE:
            return
        
        if fig is None and self.figures:
            fig = self.figures[-1]
        
        if fig is not None:
            fig.savefig(filename, dpi=300, bbox_inches='tight')
            logger.info(f"Figure saved to {filename}")

    def show(self, block: bool = True) -> None:
        """Show all figures."""
        if not MATPLOTLIB_AVAILABLE:
            return
        
        plt.show(block=block)


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Phase Portraits")
    print("="*60)
    
    pp = PhasePortraits()
    
    # Generate sample data (limit cycle)
    t = np.linspace(0, 50, 1000)
    x = 2.0 + 0.5 * np.sin(0.5 * t)
    y = 3.0 + 0.3 * np.cos(0.5 * t)
    z = 0.08 + 0.01 * np.sin(0.3 * t)
    
    # Create 2D portrait
    if MATPLOTLIB_AVAILABLE:
        fig = pp.create_2d_portrait(x, y, xlabel='V', ylabel='L', 
                                    title='V-L Phase Portrait')
        print("\nCreated 2D phase portrait")
        
        # Create Poincaré section
        pp.create_poincare_section(x, y, threshold=2.0, plane='x')
        print("Created Poincaré section")
    
    print("\n" + "="*60)
