"""
Visualization Module - 3D trajectory and sensor plotting
"""

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from typing import List, Dict, Optional, Tuple
import io
import base64


class Visualization:
    """
    Visualization tools for HADEX data
    
    Creates plots of trajectories, sensor data, and environmental maps
    """
    
    def __init__(self, style: str = 'dark_background'):
        """
        Initialize visualization
        
        Args:
            style: Matplotlib style
        """
        plt.style.use(style)
        self.figures = []
    
    def plot_trajectory(self,
                       positions: np.ndarray,
                       colors: Optional[List[str]] = None,
                       title: str = "Vehicle Trajectory",
                       show: bool = False) -> bytes:
        """
        Plot 3D vehicle trajectory
        
        Args:
            positions: Array of positions [n_steps, 3]
            colors: Colors for different segments
            title: Plot title
            
        Returns:
            PNG image as bytes
        """
        fig = plt.figure(figsize=(12, 8))
        ax = fig.add_subplot(111, projection='3d')
        
        if colors is None:
            colors = ['blue'] * len(positions)
        
        # Plot trajectory
        ax.plot(positions[:, 0], positions[:, 1], -positions[:, 2], 
               color='cyan', linewidth=2, alpha=0.8)
        
        # Mark start and end
        ax.scatter(positions[0, 0], positions[0, 1], -positions[0, 2],
                  color='green', s=100, label='Start')
        ax.scatter(positions[-1, 0], positions[-1, 1], -positions[-1, 2],
                  color='red', s=100, label='End')
        
        ax.set_xlabel('X (m)')
        ax.set_ylabel('Y (m)')
        ax.set_zlabel('Depth (m)')
        ax.set_title(title)
        ax.legend()
        
        # Equal aspect ratio
        max_range = np.max([
            np.ptp(positions[:, 0]),
            np.ptp(positions[:, 1]),
            np.ptp(positions[:, 2])
        ]) / 2.0
        
        mid_x = np.mean(positions[:, 0])
        mid_y = np.mean(positions[:, 1])
        mid_z = -np.mean(positions[:, 2])
        
        ax.set_xlim(mid_x - max_range, mid_x + max_range)
        ax.set_ylim(mid_y - max_range, mid_y + max_range)
        ax.set_zlim(mid_z - max_range, mid_z + max_range)
        
        self.figures.append(fig)
        
        if show:
            plt.show()
        
        # Convert to bytes
        buf = io.BytesIO()
        fig.savefig(buf, format='png', dpi=150, bbox_inches='tight')
        buf.seek(0)
        plt.close(fig)
        
        return buf.getvalue()
    
    def plot_oxygen_flux(self,
                        grid: np.ndarray,
                        resolution: float = 50.0,
                        colormap: str = 'viridis',
                        show: bool = False) -> bytes:
        """
        Plot oxygen flux map
        
        Args:
            grid: 2D grid of oxygen measurements
            resolution: Grid resolution in meters
            colormap: Matplotlib colormap
            
        Returns:
            PNG image as bytes
        """
        fig, ax = plt.subplots(figsize=(10, 8))
        
        im = ax.imshow(grid.T, origin='lower', cmap=colormap,
                      extent=[0, grid.shape[0]*resolution, 
                             0, grid.shape[1]*resolution])
        
        plt.colorbar(im, ax=ax, label='Oxygen Concentration (μM)')
        
        ax.set_xlabel('X (m)')
        ax.set_ylabel('Y (m)')
        ax.set_title('Benthic Oxygen Flux Map')
        
        self.figures.append(fig)
        
        if show:
            plt.show()
        
        buf = io.BytesIO()
        fig.savefig(buf, format='png', dpi=150, bbox_inches='tight')
        buf.seek(0)
        plt.close(fig)
        
        return buf.getvalue()
    
    def plot_pressure_timeseries(self,
                                time: np.ndarray,
                                pressure: np.ndarray,
                                fluctuations: Optional[List] = None,
                                show: bool = False) -> bytes:
        """
        Plot pressure time series
        
        Args:
            time: Time array (hours)
            pressure: Pressure array (MPa)
            fluctuations: Detected fluctuation events
            
        Returns:
            PNG image as bytes
        """
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8))
        
        # Main pressure plot
        ax1.plot(time, pressure, 'b-', linewidth=1, alpha=0.7)
        ax1.set_xlabel('Time (hours)')
        ax1.set_ylabel('Pressure (MPa)')
        ax1.set_title('Ambient Pressure vs Time')
        ax1.grid(True, alpha=0.3)
        
        # Detrended fluctuations
        pressure_detrend = pressure - np.mean(pressure)
        ax2.plot(time, pressure_detrend * 1000, 'r-', linewidth=1, alpha=0.7)  # Convert to kPa
        ax2.set_xlabel('Time (hours)')
        ax2.set_ylabel('Pressure Fluctuation (kPa)')
        ax2.set_title('Pressure Fluctuations')
        ax2.grid(True, alpha=0.3)
        
        # Mark detected events
        if fluctuations:
            for f in fluctuations:
                ax2.axvline(f['time'], color='yellow', alpha=0.5, linestyle='--')
        
        plt.tight_layout()
        
        self.figures.append(fig)
        
        if show:
            plt.show()
        
        buf = io.BytesIO()
        fig.savefig(buf, format='png', dpi=150, bbox_inches='tight')
        buf.seek(0)
        plt.close(fig)
        
        return buf.getvalue()
    
    def plot_bioluminescence_events(self,
                                   events: List[Dict],
                                   show: bool = False) -> bytes:
        """
        Plot bioluminescence events
        
        Args:
            events: List of bioluminescence events
            
        Returns:
            PNG image as bytes
        """
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))
        
        if events:
            # Extract data
            wavelengths = [e['wavelength'] for e in events]
            intensities = [e['intensity'] for e in events]
            durations = [e['duration'] for e in events]
            
            # Wavelength histogram
            ax1.hist(wavelengths, bins=30, color='blue', alpha=0.7, edgecolor='white')
            ax1.set_xlabel('Wavelength (nm)')
            ax1.set_ylabel('Count')
            ax1.set_title('Bioluminescence Wavelength Distribution')
            ax1.axvline(472, color='cyan', linestyle='--', label='Peak (472 nm)')
            ax1.legend()
            
            # Intensity vs Duration scatter
            sc = ax2.scatter(durations, intensities, c=wavelengths, 
                           cmap='viridis', alpha=0.6, s=50)
            ax2.set_xlabel('Duration (s)')
            ax2.set_ylabel('Intensity (photons/cm²/s)')
            ax2.set_title('Event Characteristics')
            ax2.set_yscale('log')
            plt.colorbar(sc, ax=ax2, label='Wavelength (nm)')
        else:
            ax1.text(0.5, 0.5, 'No events recorded', 
                    ha='center', va='center', transform=ax1.transAxes)
            ax2.text(0.5, 0.5, 'No events recorded', 
                    ha='center', va='center', transform=ax2.transAxes)
        
        plt.tight_layout()
        
        self.figures.append(fig)
        
        if show:
            plt.show()
        
        buf = io.BytesIO()
        fig.savefig(buf, format='png', dpi=150, bbox_inches='tight')
        buf.seek(0)
        plt.close(fig)
        
        return buf.getvalue()
    
    def plot_swarm_positions(self,
                            positions: Dict[int, np.ndarray],
                            show: bool = False) -> bytes:
        """
        Plot current swarm positions
        
        Args:
            positions: Dict of {vehicle_id: position}
            
        Returns:
            PNG image as bytes
        """
        fig = plt.figure(figsize=(12, 8))
        ax = fig.add_subplot(111, projection='3d')
        
        colors = plt.cm.rainbow(np.linspace(0, 1, len(positions)))
        
        for (vid, pos), color in zip(positions.items(), colors):
            ax.scatter(pos[0], pos[1], -pos[2], 
                      color=color, s=100, label=f'V{vid}')
        
        ax.set_xlabel('X (m)')
        ax.set_ylabel('Y (m)')
        ax.set_zlabel('Depth (m)')
        ax.set_title('Swarm Positions')
        ax.legend(loc='upper right', bbox_to_anchor=(1.15, 1))
        
        self.figures.append(fig)
        
        if show:
            plt.show()
        
        buf = io.BytesIO()
        fig.savefig(buf, format='png', dpi=150, bbox_inches='tight')
        buf.seek(0)
        plt.close(fig)
        
        return buf.getvalue()
    
    def plot_energy_harvesting(self,
                              time: np.ndarray,
                              power: np.ndarray,
                              show: bool = False) -> bytes:
        """
        Plot energy harvesting data
        
        Args:
            time: Time array (hours)
            power: Power array (mW)
            
        Returns:
            PNG image as bytes
        """
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8))
        
        # Power time series
        ax1.plot(time, power, 'g-', linewidth=1, alpha=0.7)
        ax1.axhline(np.mean(power), color='yellow', linestyle='--', 
                   label=f'Mean: {np.mean(power):.2f} mW')
        ax1.set_xlabel('Time (hours)')
        ax1.set_ylabel('Power (mW)')
        ax1.set_title('Harvested Power vs Time')
        ax1.legend()
        ax1.grid(True, alpha=0.3)
        
        # Power histogram
        ax2.hist(power, bins=50, color='green', alpha=0.7, edgecolor='white')
        ax2.set_xlabel('Power (mW)')
        ax2.set_ylabel('Count')
        ax2.set_title('Power Distribution')
        
        plt.tight_layout()
        
        self.figures.append(fig)
        
        if show:
            plt.show()
        
        buf = io.BytesIO()
        fig.savefig(buf, format='png', dpi=150, bbox_inches='tight')
        buf.seek(0)
        plt.close(fig)
        
        return buf.getvalue()
    
    def animate_events(self,
                      events: List[Dict],
                      timespan: float = 56,
                      output: str = 'events.mp4'):
        """
        Create animation of bioluminescence events
        
        Args:
            events: List of events
            timespan: Time span in hours
            output: Output filename
        """
        # This would create a video animation
        # Simplified for now
        pass
    
    def clear(self):
        """Clear all figures"""
        plt.close('all')
        self.figures = []
    
    def __repr__(self) -> str:
        return f"Visualization(figures={len(self.figures)})"
