"""
Stability maps for cave passage visualization.
Creates 2D and 3D maps of stability indices across cave systems.

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
    from matplotlib.colors import LinearSegmentedColormap
    MATPLOTLIB_AVAILABLE = True
except ImportError:
    MATPLOTLIB_AVAILABLE = False
    logger.warning("Matplotlib not available - stability maps limited")

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


class StabilityMaps:
    """
    Creates stability maps for cave passage visualization.
    """
    
    # Custom colormap for stability (red = unstable, green = stable)
    STABILITY_CMAP = LinearSegmentedColormap.from_list(
        'stability',
        ['darkred', 'red', 'orange', 'yellow', 'lightgreen', 'green'],
        N=100
    ) if MATPLOTLIB_AVAILABLE else None
    
    def __init__(self):
        """Initialize stability maps."""
        self.maps = []
        logger.info("Initialized StabilityMaps")

    def create_2d_stability_map(self,
                               x: np.ndarray,
                               y: np.ndarray,
                               stability: np.ndarray,
                               title: str = "2D Stability Map",
                               show_points: bool = True,
                               ax=None) -> Optional[plt.Figure]:
        """
        Create 2D stability map with interpolation.
        
        Args:
            x, y: Point coordinates
            stability: Stability values at points
            title: Plot title
            show_points: Show measurement points
            ax: Matplotlib axis
            
        Returns:
            Matplotlib figure or None
        """
        if not MATPLOTLIB_AVAILABLE:
            return None
        
        if ax is None:
            fig, ax = plt.subplots(figsize=(12, 10))
        else:
            fig = ax.figure
        
        # Create grid for interpolation
        xi = np.linspace(min(x), max(x), 100)
        yi = np.linspace(min(y), max(y), 100)
        Xi, Yi = np.meshgrid(xi, yi)
        
        # Simple inverse distance weighting interpolation
        Zi = self._idw_interpolation(x, y, stability, Xi, Yi)
        
        # Plot interpolated map
        contour = ax.contourf(Xi, Yi, Zi, levels=20, cmap=self.STABILITY_CMAP, alpha=0.8)
        plt.colorbar(contour, ax=ax, label='Stability Score')
        
        if show_points:
            scatter = ax.scatter(x, y, c=stability, s=50, 
                                cmap=self.STABILITY_CMAP, 
                                edgecolors='black', linewidth=1,
                                vmin=0, vmax=100)
        
        # Add contour lines
        ax.contour(Xi, Yi, Zi, levels=10, colors='black', linewidths=0.5, alpha=0.3)
        
        ax.set_xlabel('X (m)')
        ax.set_ylabel('Y (m)')
        ax.set_title(title)
        ax.set_aspect('equal')
        ax.grid(True, alpha=0.3)
        
        self.maps.append(fig)
        return fig

    def _idw_interpolation(self,
                          x: np.ndarray,
                          y: np.ndarray,
                          z: np.ndarray,
                          Xi: np.ndarray,
                          Yi: np.ndarray,
                          power: int = 2) -> np.ndarray:
        """Inverse distance weighting interpolation."""
        Zi = np.zeros_like(Xi)
        
        for i in range(Xi.shape[0]):
            for j in range(Xi.shape[1]):
                # Calculate distances to all points
                distances = np.sqrt((Xi[i, j] - x)**2 + (Yi[i, j] - y)**2)
                
                # Avoid division by zero
                distances = np.maximum(distances, 1e-10)
                
                # Calculate weights
                weights = 1.0 / (distances ** power)
                
                # Weighted average
                Zi[i, j] = np.sum(weights * z) / np.sum(weights)
        
        return Zi

    def create_3d_stability_map(self,
                               x: np.ndarray,
                               y: np.ndarray,
                               z: np.ndarray,
                               stability: np.ndarray,
                               title: str = "3D Stability Map") -> Optional[go.Figure]:
        """
        Create 3D stability map using Plotly.
        
        Args:
            x, y, z: Point coordinates
            stability: Stability values
            title: Plot title
            
        Returns:
            Plotly figure or None
        """
        if not PLOTLY_AVAILABLE:
            return None
        
        # Create color scale
        colorscale = [
            [0, 'darkred'],
            [0.25, 'red'],
            [0.5, 'orange'],
            [0.75, 'yellow'],
            [1, 'green']
        ]
        
        fig = go.Figure(data=[
            go.Scatter3d(
                x=x, y=y, z=z,
                mode='markers',
                marker=dict(
                    size=5,
                    color=stability,
                    colorscale=colorscale,
                    cmin=0,
                    cmax=100,
                    showscale=True,
                    colorbar=dict(title="Stability")
                ),
                text=[f"Stability: {s:.1f}" for s in stability]
            )
        ])
        
        # Add surface interpolation (simplified - just for visualization)
        try:
            from scipy.interpolate import griddata
            
            # Create grid
            xi = np.linspace(min(x), max(x), 30)
            yi = np.linspace(min(y), max(y), 30)
            Xi, Yi = np.meshgrid(xi, yi)
            
            # Interpolate stability
            Zi = griddata((x, y), stability, (Xi, Yi), method='cubic')
            
            # Add surface
            fig.add_trace(go.Surface(
                x=xi, y=yi, z=np.zeros_like(Xi) + min(z),
                surfacecolor=Zi,
                colorscale=colorscale,
                showscale=False,
                opacity=0.5
            ))
        except:
            pass
        
        fig.update_layout(
            title=title,
            scene=dict(
                xaxis_title='X (m)',
                yaxis_title='Y (m)',
                zaxis_title='Depth (m)'
            ),
            height=600
        )
        
        return fig

    def create_stability_profile(self,
                                distance: np.ndarray,
                                stability: np.ndarray,
                                title: str = "Stability Profile",
                                threshold: float = 60.0) -> Optional[plt.Figure]:
        """
        Create 2D stability profile along passage.
        
        Args:
            distance: Distance along passage (m)
            stability: Stability values
            title: Plot title
            threshold: Stability threshold
            
        Returns:
            Matplotlib figure or None
        """
        if not MATPLOTLIB_AVAILABLE:
            return None
        
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8), 
                                       gridspec_kw={'height_ratios': [3, 1]})
        
        # Main profile
        ax1.plot(distance, stability, 'b-', linewidth=2, label='Stability')
        ax1.axhline(y=threshold, color='r', linestyle='--', 
                   label=f'Threshold ({threshold})')
        
        # Color zones based on stability
        ax1.fill_between(distance, 0, stability, 
                         where=(stability >= threshold),
                         color='green', alpha=0.3, label='Safe')
        ax1.fill_between(distance, 0, stability,
                         where=(stability < threshold),
                         color='red', alpha=0.3, label='Unsafe')
        
        ax1.set_xlabel('Distance along passage (m)')
        ax1.set_ylabel('Stability Score')
        ax1.set_title(title)
        ax1.set_ylim(0, 100)
        ax1.grid(True, alpha=0.3)
        ax1.legend()
        
        # Risk zones
        risk_levels = []
        for s in stability:
            if s >= 80:
                risk_levels.append(1)  # Safe
            elif s >= 60:
                risk_levels.append(2)  # Caution
            elif s >= 40:
                risk_levels.append(3)  # Unstable
            elif s >= 20:
                risk_levels.append(4)  # High Risk
            else:
                risk_levels.append(5)  # Critical
        
        colors = ['green', 'yellow', 'orange', 'red', 'darkred']
        for i in range(len(distance) - 1):
            ax2.barh(0, distance[i+1] - distance[i], 
                    left=distance[i], height=0.5,
                    color=colors[risk_levels[i]-1],
                    edgecolor='black', linewidth=0.5)
        
        ax2.set_xlabel('Distance along passage (m)')
        ax2.set_yticks([])
        ax2.set_title('Risk Zones')
        ax2.set_xlim(ax1.get_xlim())
        
        # Add legend for risk zones
        legend_elements = []
        for level, color in zip(['Safe', 'Caution', 'Unstable', 'High Risk', 'Critical'],
                                colors):
            legend_elements.append(plt.Rectangle((0, 0), 1, 1, facecolor=color))
        ax2.legend(legend_elements, ['Safe', 'Caution', 'Unstable', 'High Risk', 'Critical'],
                  loc='upper right', ncol=5)
        
        plt.tight_layout()
        
        self.maps.append(fig)
        return fig

    def create_risk_heatmap(self,
                           x: np.ndarray,
                           y: np.ndarray,
                           risk_scores: Dict[str, np.ndarray],
                           title: str = "Risk Component Heatmap") -> Optional[plt.Figure]:
        """
        Create heatmap of multiple risk components.
        
        Args:
            x, y: Grid coordinates
            risk_scores: Dictionary of risk component name to scores
            title: Plot title
            
        Returns:
            Matplotlib figure or None
        """
        if not MATPLOTLIB_AVAILABLE:
            return None
        
        n_components = len(risk_scores)
        fig, axes = plt.subplots(1, n_components, figsize=(5*n_components, 4))
        
        if n_components == 1:
            axes = [axes]
        
        for ax, (name, scores) in zip(axes, risk_scores.items()):
            # Reshape scores to grid
            n_x = len(np.unique(x))
            n_y = len(np.unique(y))
            scores_grid = scores.reshape(n_y, n_x)
            
            im = ax.imshow(scores_grid, cmap='RdYlGn_r', aspect='auto',
                          extent=[min(x), max(x), min(y), max(y)],
                          origin='lower', vmin=0, vmax=100)
            plt.colorbar(im, ax=ax, label='Score')
            
            ax.set_xlabel('X (m)')
            ax.set_ylabel('Y (m)')
            ax.set_title(name)
        
        plt.suptitle(title)
        plt.tight_layout()
        
        self.maps.append(fig)
        return fig

    def create_stability_animation(self,
                                  x: np.ndarray,
                                  y: np.ndarray,
                                  stability_series: List[np.ndarray],
                                  times: List[float],
                                  filename: str = "stability_animation.gif") -> None:
        """
        Create animation of stability changes over time.
        
        Args:
            x, y: Point coordinates
            stability_series: List of stability arrays at each time
            times: Corresponding times
            filename: Output filename
        """
        if not MATPLOTLIB_AVAILABLE:
            return
        
        from matplotlib.animation import FuncAnimation
        
        fig, ax = plt.subplots(figsize=(10, 8))
        
        def animate(frame):
            ax.clear()
            stability = stability_series[frame]
            
            # Create scatter plot
            scatter = ax.scatter(x, y, c=stability, s=50,
                                cmap=self.STABILITY_CMAP,
                                vmin=0, vmax=100,
                                edgecolors='black', linewidth=1)
            
            ax.set_xlabel('X (m)')
            ax.set_ylabel('Y (m)')
            ax.set_title(f'Stability Map at t = {times[frame]:.1f}')
            ax.set_aspect('equal')
            ax.grid(True, alpha=0.3)
            
            return scatter,
        
        anim = FuncAnimation(fig, animate, frames=len(stability_series), 
                           interval=200, blit=True)
        anim.save(filename, writer='pillow', fps=5)
        logger.info(f"Animation saved to {filename}")

    def save_map(self, filename: str, fig=None) -> None:
        """
        Save map 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.maps:
            fig = self.maps[-1]
        
        if fig is not None:
            fig.savefig(filename, dpi=300, bbox_inches='tight')
            logger.info(f"Map saved to {filename}")


# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    print("="*60)
    print("CAVORA - Stability Maps")
    print("="*60)
    
    sm = StabilityMaps()
    
    # Generate sample data
    np.random.seed(42)
    n_points = 50
    x = np.random.uniform(0, 100, n_points)
    y = np.random.uniform(0, 100, n_points)
    stability = np.random.uniform(30, 100, n_points)
    
    if MATPLOTLIB_AVAILABLE:
        # Create 2D stability map
        fig = sm.create_2d_stability_map(x, y, stability, title="Sample Stability Map")
        print("\nCreated 2D stability map")
        
        # Create stability profile
        distance = np.sort(np.random.uniform(0, 500, 100))
        profile_stability = 80 - 20 * np.sin(distance/50) + np.random.randn(100)*5
        profile_stability = np.clip(profile_stability, 0, 100)
        
        fig2 = sm.create_stability_profile(distance, profile_stability)
        print("Created stability profile")
    
    print("\n" + "="*60)
