#!/usr/bin/env python3
"""
3D Array Visualizer with Interactive Controls
===============================================

This tool visualizes slices of 3D arrays with:
- Interactive slider to select different depth slices
- Play/Pause button to animate through slices
- Real-time visualization updates
- Custom colormap support
- Stacked plot showing evolution of health states

Features:
- Slider control for navigating through 3D array slices
- Play button to automatically cycle through slices
- Color-coded visualization with customizable colormaps
- Sample data generation for immediate demonstration
- Stacked plot showing evolution of health states (healthy, immune, infected, dead)

Usage:
Run the script to see the visualization interface:
    python3 3d_visualizer_tool.py

Controls:
- Use the slider to select specific slice
- Click "Play" to start animation
- Click "Pause" to stop animation

To load a specific numpy array file:
    python3 3d_visualizer_tool.py --file evolution_map.npy
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap, ListedColormap
from matplotlib.widgets import Slider, Button
import matplotlib.animation as animation
import argparse
import sys

def get_entity_values():
    return {
        'empty': 0,
        'water': 1,
        'medical facility': 2,
        'healthy': 5,
        'vaccinated': 6,
        'infected': 10,
        'dead': 20
    }

def calculate_health_stats(data_3d, ent2ind):
    """
    Calculate health statistics for each slice of the 3D data.
    Parameters:
    data_3d: numpy array of shape (depth, height, width)
    
    Returns:
    stats: list of dictionaries containing health counts for each slice
    """
    stats = []
    for slice_idx in range(data_3d.shape[0]):
        slice_data = data_3d[slice_idx]
        
        # Count different health states
        empty_count = np.sum(slice_data == ent2ind["empty"])
        water_count = np.sum(slice_data == ent2ind["water"])
        medical_count = np.sum(slice_data == ent2ind["medical facility"])
        healthy_count = np.sum(slice_data == ent2ind["healthy"])
        immune_count = np.sum(slice_data == ent2ind["vaccinated"])
        infected_count = np.sum(slice_data == ent2ind["infected"])
        dead_count = np.sum(slice_data == ent2ind["dead"])
        
        stats.append({
            'slice': slice_idx,
            'empty': empty_count,
            'water': water_count,
            'medical': medical_count,
            'healthy': healthy_count,
            'immune': immune_count,
            'infected': infected_count,
            'dead': dead_count
        })
    
    return stats

def inverse_dictionary(d):
    rd = {}
    for k,v in d.items():
        rd[v] = k
    return rd

def get_evolution_colors():
    """
    Get the color mapping for health states.
    """
    # Default colors for health states
    colors = {
        'empty': [1, 1, 1, 1],      # white for empty space
        'water': [0, 0, 1, 1],      # blue for water
        'medical facility': [1, 0, 1, 1],      # magenta for medical facility
        'healthy': [0.5, 0.5, 0.5, 1], # grey for healthy human
        'vaccinated': [0, 1, 0, 1],      # green for vaccinated human
        'infected': [1, 0, 0, 1],     # red for infected human
        'dead': [0, 0, 0, 1]      # black for dead human
    }
    
    # If we have a custom colormap, we'll use it for the visualization
    # but still return the color mapping for legend purposes
    return colors

def evolution_colormap(ind2entity_val):
    # Define color mappings (RGBA values)
    colors = get_evolution_colors()
    entitis = get_entity_values()
    entity_val2name = inverse_dictionary(entitis)

    # Create a ListedColormap for proper matplotlib usage
    color_list = [colors[entity_val2name[ind2entity_val[i]]] for i in ind2entity_val.keys()]
    cmap = ListedColormap(color_list)
    return cmap


def visualize_3d_array(data_3d_in):
    """
    Visualize a 3D array with interactive controls and health evolution plot
    
    Parameters:
    data_3d: numpy array of shape (depth, height, width)
    colormap: matplotlib colormap to use (optional)
    """
    
    # entity can arbitrary numerical label here we compactify this notation
    vals, fi, ri= np.unique(data_3d_in, return_index=True, return_inverse=True)
    data_3d = ri.reshape(data_3d_in.shape)
    ind2entity_val={}
    for i in range(len(fi)):
        ind2entity_val[i] = int(data_3d_in.flatten()[ fi[i] ])
    val2ind = inverse_dictionary(ind2entity_val)
    entities = get_entity_values()
    ent2ind = {}
    for k, v in entities.items():
        ent2ind[k] = val2ind[v]

    # Setup the plot
    fig = plt.figure(figsize=(14, 12))
    
    # Create subplots: main visualization and health evolution
    gs = fig.add_gridspec(3, 1, height_ratios=[2, 1, 0.2], hspace=0.3)
    
    # Main visualization subplot
    ax_main = fig.add_subplot(gs[0, 0])
    
    # Health evolution subplot
    ax_health = fig.add_subplot(gs[1, 0])
    
    # Animation controls subplot
    ax_controls = fig.add_subplot(gs[2, 0])
    
    # Store data and parameters
    depth, height, width = data_3d.shape
    
    # Initial slice
    current_slice = depth // 2
    
    # Create our custom colormap
    cmap_used = evolution_colormap(ind2entity_val)
    
    # Display initial slice
    im = ax_main.imshow(data_3d[current_slice], cmap=cmap_used, vmin=data_3d.min(), vmax=data_3d.max(), aspect='auto')
    ax_main.set_title(f'Slice {current_slice} of {depth} - Depth Index: {current_slice}')
    ax_main.set_xlabel('Width')
    ax_main.set_ylabel('Height')
    
    # Get color mapping for legend
    health_colors = get_evolution_colors()
    
    # Remove colorbar and add legend instead
    # Create legend under the main image with colors that represent health states intuitively
    # These colors are chosen to be visually distinct and meaningful for health states
    legend_elements = [
        plt.Rectangle((0,0),1,1, facecolor=health_colors['medical facility'], alpha=0.8),    # medical (2)
        plt.Rectangle((0,0),1,1, facecolor=health_colors['healthy'], alpha=0.8),   # healthy (5)
        plt.Rectangle((0,0),1,1, facecolor=health_colors['vaccinated'], alpha=0.8),   # immune (6)
        plt.Rectangle((0,0),1,1, facecolor=health_colors['infected'], alpha=0.8), # infected (10)
        plt.Rectangle((0,0),1,1, facecolor=health_colors['dead'], alpha=0.8)   # dead (20)
    ]
    legend_labels = ['Medical Facilities', 'Healthy', 'Immune', 'Infected', 'Dead']
    
    # Add legend under the main image
    ax_main.legend(legend_elements, legend_labels, loc='lower center', bbox_to_anchor=(0.5, -0.1), 
                  ncol=5, frameon=True, fancybox=True, shadow=True)
    
    # Calculate health statistics for all slices
    health_stats = calculate_health_stats(data_3d, ent2ind)
    
    # Plot health evolution
    slice_indices = [stat['slice'] for stat in health_stats]
    empty_counts = [stat['empty'] for stat in health_stats]
    water_counts = [stat['water'] for stat in health_stats]
    medical_counts = [stat['medical'] for stat in health_stats]
    healthy_counts = [stat['healthy'] for stat in health_stats]
    immune_counts = [stat['immune'] for stat in health_stats]
    infected_counts = [stat['infected'] for stat in health_stats]
    dead_counts = [stat['dead'] for stat in health_stats]
    
    # Stack the data for the stacked plot
    stacked_data = np.column_stack([empty_counts, water_counts, medical_counts, healthy_counts, immune_counts, infected_counts, dead_counts])
    
    # Create stacked area plot with appropriate colors
    ax_health.stackplot(slice_indices, stacked_data.T, 
                        labels=['Empty', 'Water', 'Medical Facilities', 'Healthy', 'Immune', 'Infected', 'Dead'], 
                        colors=[health_colors['empty'], health_colors['water'], health_colors['medical facility'], health_colors['healthy'], health_colors['vaccinated'], health_colors['infected'], health_colors['dead']])
    ax_health.set_title('Evolution of Population Health States')
    ax_health.set_xlabel('Time Slice')
    ax_health.set_ylabel('Population Count')
    ax_health.legend(loc='upper right')
    ax_health.grid(True, alpha=0.3)
    
    # Create slider
    ax_slider = plt.axes([0.2, 0.15, 0.5, 0.03])
    slider = Slider(
        ax=ax_slider,
        label='Slice Index',
        valmin=0,
        valmax=depth - 1,
        valinit=current_slice,
        valfmt='%d'
    )
    
    # Create single button for play/pause in the controls area
    play_pause_button = Button(ax_controls, 'Play')
    # Position the button vertically centered in the controls area
    play_pause_button.ax.set_position([0.45, 0.1, 0.1, 0.04])
    
    # Animation state
    anim = None
    is_playing = False
    
    def update_slice(val):
        """Update displayed slice"""
        nonlocal current_slice
        current_slice = int(val)
        im.set_array(data_3d[current_slice])
        ax_main.set_title(f'Slice {current_slice} of {depth} - Depth Index: {current_slice}')
        
        # Update the health evolution plot to highlight current slice
        ax_health.clear()
        ax_health.stackplot(slice_indices, stacked_data.T, 
                            labels=['Empty', 'Water', 'Medical Facilities', 'Healthy', 'Immune', 'Infected', 'Dead'], 
                            colors=[health_colors['empty'], health_colors['water'], health_colors['medical facility'], health_colors['healthy'], health_colors['vaccinated'], health_colors['infected'], health_colors['dead']])
        ax_health.set_title('Evolution of Population Health States')
        ax_health.set_xlabel('Time Slice')
        ax_health.set_ylabel('Population Count')
        ax_health.legend(loc='upper right')
        ax_health.grid(True, alpha=0.3)
        ax_health.axvline(x=current_slice, color='red', linestyle='--', alpha=0.7)
        fig.canvas.draw_idle()
    
    def play_pause(event):
        """Toggle play/pause animation"""
        nonlocal is_playing, anim, current_slice
        if not is_playing:
            # Start playing
            is_playing = True
            play_pause_button.label.set_text('Pause')
            # Create animation if it doesn't exist
            if anim is None:
                anim = animation.FuncAnimation(
                    fig, animate, interval=300, blit=False, cache_frame_data=False
                )
            else:
                try:
                    anim.event_source.start()
                except:
                    # If animation exists but stopped, just restart it
                    anim.event_source.start()
        else:
            # Pause playing
            is_playing = False
            play_pause_button.label.set_text('Play')
            if anim is not None:
                try:
                    anim.event_source.stop()
                except:
                    pass
    
    def animate(frame):
        """Animation function"""
        nonlocal current_slice, is_playing
        if is_playing:
            current_slice = (current_slice + 1) % depth
            # Update slider and display - ensure proper update
            slider.set_val(current_slice)
            im.set_array(data_3d[current_slice])
            ax_main.set_title(f'Slice {current_slice} of {depth} - Depth Index: {current_slice}')
            
            # Update the health evolution plot to highlight current slice
            ax_health.clear()
            ax_health.stackplot(slice_indices, stacked_data.T, 
                                labels=['Empty', 'Water', 'Medical Facilities', 'Healthy', 'Immune', 'Infected', 'Dead'], 
                                colors=[health_colors['empty'], health_colors['water'], health_colors['medical facility'], health_colors['healthy'], health_colors['vaccinated'], health_colors['infected'], health_colors['dead']])
            ax_health.set_title('Evolution of Population Health States')
            ax_health.set_xlabel('Time Slice')
            ax_health.set_ylabel('Population Count')
            ax_health.legend(loc='upper right')
            ax_health.grid(True, alpha=0.3)
            ax_health.axvline(x=current_slice, color='red', linestyle='--', alpha=0.7)
            
            # Force canvas update
            fig.canvas.draw_idle()
        return [im]
    
    # Connect events
    slider.on_changed(update_slice)
    play_pause_button.on_clicked(play_pause)
    
    print(f"Visualizing 3D array with shape: {data_3d.shape}")
    print("Controls:")
    print("- Drag slider to select specific slice")
    print("- Click 'Play' to start animation")
    print("- Click 'Pause' to stop animation")
    print("\nClose the visualization window to exit.")
    
    # Show the visualization
    plt.show()


def main():
    parser = argparse.ArgumentParser(description='3D Array Visualizer')
    parser.add_argument('--file', '-f', type=str, help='Path to numpy array file (.npy) to visualize', required=True)
    
    args = parser.parse_args()
    
    print("3D Array Visualizer")
    print("=" * 30)
    
    if args.file:
        # Load data from numpy file
        try:
            data = np.load(args.file)
            print(f"Loaded 3D array from '{args.file}' with shape: {data.shape}")
        except Exception as e:
            print(f"Error loading file '{args.file}': {e}")
            sys.exit(1)
    
    visualize_3d_array(data)


if __name__ == "__main__":
    main()