"""
Network Visualizer

Advanced visualization tools for neural network architectures,
training progress, and model interpretability.
"""

import numpy as np
import pandas as pd
from typing import Dict, List, Any, Optional, Tuple
import matplotlib.pyplot as plt
import seaborn as sns
import networkx as nx
from matplotlib.patches import Rectangle, FancyBboxPatch
import warnings
from dataclasses import dataclass

# Try to import additional visualization libraries
try:
    import plotly.graph_objects as go
    import plotly.express as px
    from plotly.subplots import make_subplots
    PLOTLY_AVAILABLE = True
except ImportError:
    PLOTLY_AVAILABLE = False

try:
    from sklearn.manifold import TSNE
    from sklearn.decomposition import PCA
    SKLEARN_AVAILABLE = True
except ImportError:
    SKLEARN_AVAILABLE = False


@dataclass
class VisualizationConfig:
    """Configuration for visualization settings"""
    figure_size: Tuple[int, int] = (12, 8)
    dpi: int = 100
    style: str = 'seaborn-v0_8'
    color_palette: str = 'viridis'
    save_format: str = 'png'
    interactive: bool = True


class NetworkVisualizer:
    """
    Comprehensive neural network visualization toolkit with support for
    architecture diagrams, training progress, and model interpretability.
    """
    
    def __init__(self, config: Optional[VisualizationConfig] = None):
        """Initialize the visualizer with configuration"""
        self.config = config or VisualizationConfig()
        self.setup_style()
        
    def setup_style(self):
        """Setup matplotlib style"""
        try:
            plt.style.use(self.config.style)
        except:
            plt.style.use('default')
        
        sns.set_palette(self.config.color_palette)
    
    def plot_network_architecture(self, layers: List[Dict[str, Any]], 
                                save_path: Optional[str] = None,
                                interactive: Optional[bool] = None) -> None:
        """
        Visualize neural network architecture
        
        Args:
            layers: List of layer configurations
            save_path: Path to save the visualization
            interactive: Whether to create interactive plot
        """
        interactive = interactive if interactive is not None else self.config.interactive
        
        if interactive and PLOTLY_AVAILABLE:
            self._plot_network_architecture_plotly(layers, save_path)
        else:
            self._plot_network_architecture_matplotlib(layers, save_path)
    
    def _plot_network_architecture_matplotlib(self, layers: List[Dict[str, Any]], 
                                           save_path: Optional[str] = None):
        """Plot network architecture using matplotlib"""
        fig, ax = plt.subplots(figsize=self.config.figure_size)
        
        # Calculate positions
        max_width = max([layer.get('units', layer.get('filters', 10)) for layer in layers])
        layer_spacing = 2.0
        neuron_spacing = 0.5
        
        positions = {}
        layer_positions = []
        
        for i, layer in enumerate(layers):
            x = i * layer_spacing
            layer_units = layer.get('units', layer.get('filters', 10))
            
            # Normalize layer width
            normalized_width = (layer_units / max_width) * 3
            
            # Position neurons in this layer
            for j in range(min(layer_units, 10)):  # Limit to 10 neurons for visualization
                y = (j - layer_units/2) * neuron_spacing
                positions[f"{i}_{j}"] = (x, y)
            
            layer_positions.append((x, normalized_width, layer))
        
        # Draw layers
        for i, (x, width, layer) in enumerate(layer_positions):
            layer_type = layer.get('type', 'dense')
            
            if layer_type in ['dense', 'Dense']:
                # Draw neurons as circles
                layer_units = layer.get('units', 10)
                for j in range(min(layer_units, 10)):
                    y = (j - layer_units/2) * 0.5
                    circle = plt.Circle((x, y), 0.1, color='skyblue', ec='black', linewidth=1)
                    ax.add_patch(circle)
                
                # Draw layer box
                rect = FancyBboxPatch(
                    (x - 0.2, -layer_units * 0.25), 0.4, layer_units * 0.5,
                    boxstyle="round,pad=0.05", 
                    facecolor='lightblue', alpha=0.3, edgecolor='blue'
                )
                ax.add_patch(rect)
                
            elif layer_type in ['conv2d', 'Conv2D']:
                # Draw as rectangle
                rect = FancyBboxPatch(
                    (x - 0.3, -1), 0.6, 2,
                    boxstyle="round,pad=0.05",
                    facecolor='lightgreen', alpha=0.3, edgecolor='green'
                )
                ax.add_patch(rect)
                ax.text(x, 0, f"Conv\n{layer.get('filters', 32)}", 
                       ha='center', va='center', fontsize=8)
                
            elif layer_type in ['dropout', 'Dropout']:
                # Draw as dashed box
                rect = FancyBboxPatch(
                    (x - 0.3, -1), 0.6, 2,
                    boxstyle="round,pad=0.05",
                    facecolor='lightyellow', alpha=0.3, edgecolor='orange',
                    linestyle='--', linewidth=2
                )
                ax.add_patch(rect)
                ax.text(x, 0, f"Dropout\n{layer.get('rate', 0.5)}", 
                       ha='center', va='center', fontsize=8)
                
            elif layer_type in ['lstm', 'LSTM']:
                # Draw as special shape
                rect = FancyBboxPatch(
                    (x - 0.3, -1), 0.6, 2,
                    boxstyle="round,pad=0.05",
                    facecolor='lightcoral', alpha=0.3, edgecolor='red'
                )
                ax.add_patch(rect)
                ax.text(x, 0, f"LSTM\n{layer.get('units', 64)}", 
                       ha='center', va='center', fontsize=8)
            
            # Add layer label
            ax.text(x, -2.5, layer.get('name', f'Layer {i}'), 
                   ha='center', va='top', fontsize=10, fontweight='bold')
        
        # Draw connections (simplified)
        for i in range(len(layer_positions) - 1):
            x1, _, _ = layer_positions[i]
            x2, _, _ = layer_positions[i + 1]
            
            # Draw a few representative connections
            for j in range(3):
                y1 = (j - 1.5) * 0.5
                y2 = (j - 1.5) * 0.5
                ax.plot([x1 + 0.2, x2 - 0.2], [y1, y2], 'gray', alpha=0.3, linewidth=0.5)
        
        ax.set_xlim(-1, len(layers) * 2)
        ax.set_ylim(-3, 3)
        ax.set_aspect('equal')
        ax.axis('off')
        ax.set_title('Neural Network Architecture', fontsize=14, fontweight='bold')
        
        plt.tight_layout()
        if save_path:
            plt.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
        plt.show()
    
    def _plot_network_architecture_plotly(self, layers: List[Dict[str, Any]], 
                                        save_path: Optional[str] = None):
        """Plot network architecture using plotly (interactive)"""
        fig = go.Figure()
        
        # Create nodes for each layer
        node_x = []
        node_y = []
        node_text = []
        node_info = []
        
        for i, layer in enumerate(layers):
            x = i * 2
            layer_type = layer.get('type', 'dense')
            
            if layer_type in ['dense', 'Dense']:
                units = layer.get('units', 10)
                # Create multiple nodes for neurons
                for j in range(min(units, 5)):  # Limit to 5 for visualization
                    y = (j - units/2) * 0.5
                    node_x.append(x)
                    node_y.append(y)
                    node_text.append(f"{layer.get('name', 'Dense')}<br>Unit {j+1}")
                    node_info.append(f"Type: {layer_type}<br>Units: {units}<br>Activation: {layer.get('activation', 'None')}")
            else:
                # Single node for other layer types
                node_x.append(x)
                node_y.append(0)
                node_text.append(layer.get('name', layer_type))
                node_info.append(f"Type: {layer_type}<br>Params: {layer}")
        
        # Add nodes
        fig.add_trace(go.Scatter(
            x=node_x, y=node_y,
            mode='markers+text',
            hoverinfo='text',
            text=node_text,
            textposition="top center",
            hovertext=node_info,
            marker=dict(
                size=20,
                color='lightblue',
                line=dict(width=2, color='darkblue')
            ),
            name='Layers'
        ))
        
        # Add connections (simplified)
        for i in range(len(layers) - 1):
            x1 = i * 2
            x2 = (i + 1) * 2
            
            fig.add_trace(go.Scatter(
                x=[x1, x2], y=[0, 0],
                mode='lines',
                line=dict(width=1, color='gray'),
                showlegend=False,
                hoverinfo='none'
            ))
        
        fig.update_layout(
            title="Neural Network Architecture (Interactive)",
            showlegend=False,
            hovermode='closest',
            margin=dict(b=20,l=5,r=5,t=40),
            annotations=[ dict(
                text="Hover over layers for details",
                showarrow=False,
                xref="paper", yref="paper",
                x=0.005, y=-0.002,
                xanchor='left', yanchor='bottom',
                font=dict(size=10)
            )],
            xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
            yaxis=dict(showgrid=False, zeroline=False, showticklabels=False)
        )
        
        if save_path:
            fig.write_html(save_path.replace('.png', '.html'))
        
        fig.show()
    
    def plot_training_progress(self, history: List[Dict[str, Any]], 
                             save_path: Optional[str] = None) -> None:
        """
        Plot training progress with multiple metrics
        
        Args:
            history: Training history dictionary
            save_path: Path to save the visualization
        """
        if not history:
            print("No training history available")
            return
        
        epochs = list(range(1, len(history) + 1))
        
        # Determine number of subplots needed
        metrics = set()
        for epoch_data in history:
            metrics.update(epoch_data.keys())
        
        # Remove common keys that aren't metrics
        metrics = metrics - {'epoch', 'timestamp'}
        
        n_metrics = len(metrics)
        if n_metrics == 0:
            print("No metrics found in training history")
            return
        
        # Create subplots
        cols = min(2, n_metrics)
        rows = (n_metrics + cols - 1) // cols
        
        fig, axes = plt.subplots(rows, cols, figsize=(cols * 6, rows * 4))
        if n_metrics == 1:
            axes = [axes]
        elif rows == 1:
            axes = axes.reshape(1, -1)
        
        for i, metric in enumerate(sorted(metrics)):
            row = i // cols
            col = i % cols
            ax = axes[row, col] if rows > 1 else axes[col]
            
            values = [epoch_data.get(metric) for epoch_data in history]
            if values and all(v is not None for v in values):
                ax.plot(epochs, values, marker='o', linewidth=2, markersize=4)
                ax.set_title(f'{metric.replace("_", " ").title()}')
                ax.set_xlabel('Epoch')
                ax.set_ylabel(metric.replace("_", " ").title())
                ax.grid(True, alpha=0.3)
        
        # Hide empty subplots
        for i in range(n_metrics, rows * cols):
            row = i // cols
            col = i % cols
            if rows > 1:
                axes[row, col].set_visible(False)
            else:
                axes[col].set_visible(False)
        
        plt.tight_layout()
        if save_path:
            plt.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
        plt.show()
    
    def plot_loss_landscape(self, history: List[Dict[str, Any]], 
                          save_path: Optional[str] = None) -> None:
        """
        Plot 3D loss landscape visualization
        
        Args:
            history: Training history
            save_path: Path to save the visualization
        """
        if not PLOTLY_AVAILABLE:
            print("Plotly is required for 3D visualization")
            return
        
        train_losses = [h.get('train_loss') for h in history]
        val_losses = [h.get('val_loss') for h in history]
        
        if not train_losses:
            print("No loss data available")
            return
        
        epochs = list(range(1, len(train_losses) + 1))
        
        fig = go.Figure(data=[go.Scatter3d(
            x=epochs,
            y=train_losses,
            z=[0] * len(epochs),
            mode='markers+lines',
            name='Training Loss',
            marker=dict(size=5, color='blue'),
            line=dict(width=3)
        )])
        
        if val_losses and all(v is not None for v in val_losses):
            fig.add_trace(go.Scatter3d(
                x=epochs,
                y=val_losses,
                z=[1] * len(epochs),
                mode='markers+lines',
                name='Validation Loss',
                marker=dict(size=5, color='red'),
                line=dict(width=3)
            ))
        
        fig.update_layout(
            title='3D Loss Landscape',
            scene=dict(
                xaxis_title='Epoch',
                yaxis_title='Loss',
                zaxis_title='Dataset'
            ),
            width=800,
            height=600
        )
        
        if save_path:
            fig.write_html(save_path.replace('.png', '.html'))
        
        fig.show()
    
    def plot_feature_importance(self, importance_scores: Dict[str, float],
                              top_n: int = 20, save_path: Optional[str] = None) -> None:
        """
        Plot feature importance scores
        
        Args:
            importance_scores: Dictionary of feature names and importance scores
            top_n: Number of top features to display
            save_path: Path to save the visualization
        """
        if not importance_scores:
            print("No importance scores provided")
            return
        
        # Sort by importance
        sorted_features = sorted(importance_scores.items(), key=lambda x: x[1], reverse=True)
        top_features = sorted_features[:top_n]
        
        features, scores = zip(*top_features)
        
        plt.figure(figsize=self.config.figure_size)
        
        # Create horizontal bar plot
        y_pos = np.arange(len(features))
        bars = plt.barh(y_pos, scores, color='skyblue', edgecolor='navy', alpha=0.7)
        
        # Add value labels on bars
        for i, (bar, score) in enumerate(zip(bars, scores)):
            plt.text(bar.get_width() + 0.001, bar.get_y() + bar.get_height()/2,
                    f'{score:.3f}', ha='left', va='center', fontsize=9)
        
        plt.yticks(y_pos, features)
        plt.xlabel('Importance Score')
        plt.title(f'Top {top_n} Feature Importance Scores')
        plt.gca().invert_yaxis()  # Highest importance at top
        plt.grid(axis='x', alpha=0.3)
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
        plt.show()
    
    def plot_confusion_matrix(self, y_true: np.ndarray, y_pred: np.ndarray,
                            class_names: Optional[List[str]] = None,
                            save_path: Optional[str] = None) -> None:
        """
        Plot confusion matrix
        
        Args:
            y_true: True labels
            y_pred: Predicted labels
            class_names: Names of classes
            save_path: Path to save the visualization
        """
        from sklearn.metrics import confusion_matrix
        
        cm = confusion_matrix(y_true, y_pred)
        
        plt.figure(figsize=self.config.figure_size)
        
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                   xticklabels=class_names, yticklabels=class_names)
        plt.title('Confusion Matrix')
        plt.xlabel('Predicted Label')
        plt.ylabel('True Label')
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
        plt.show()
    
    def plot_learning_curves(self, train_scores: List[float], val_scores: List[float],
                           metric_name: str = 'Score', save_path: Optional[str] = None) -> None:
        """
        Plot learning curves
        
        Args:
            train_scores: Training scores
            val_scores: Validation scores
            metric_name: Name of the metric
            save_path: Path to save the visualization
        """
        epochs = list(range(1, len(train_scores) + 1))
        
        plt.figure(figsize=self.config.figure_size)
        
        plt.plot(epochs, train_scores, 'b-', label=f'Training {metric_name}', linewidth=2)
        if val_scores:
            plt.plot(epochs, val_scores, 'r-', label=f'Validation {metric_name}', linewidth=2)
        
        plt.xlabel('Epoch')
        plt.ylabel(metric_name)
        plt.title(f'Learning Curves - {metric_name}')
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
        plt.show()
    
    def plot_residuals(self, y_true: np.ndarray, y_pred: np.ndarray,
                     save_path: Optional[str] = None) -> None:
        """
        Plot residual analysis for regression models
        
        Args:
            y_true: True values
            y_pred: Predicted values
            save_path: Path to save the visualization
        """
        residuals = y_true - y_pred
        
        fig, axes = plt.subplots(2, 2, figsize=(12, 10))
        
        # Residuals vs Predicted
        axes[0, 0].scatter(y_pred, residuals, alpha=0.6)
        axes[0, 0].axhline(y=0, color='r', linestyle='--')
        axes[0, 0].set_xlabel('Predicted Values')
        axes[0, 0].set_ylabel('Residuals')
        axes[0, 0].set_title('Residuals vs Predicted')
        axes[0, 0].grid(True, alpha=0.3)
        
        # Histogram of residuals
        axes[0, 1].hist(residuals, bins=30, alpha=0.7, color='skyblue', edgecolor='black')
        axes[0, 1].set_xlabel('Residuals')
        axes[0, 1].set_ylabel('Frequency')
        axes[0, 1].set_title('Distribution of Residuals')
        axes[0, 1].grid(True, alpha=0.3)
        
        # Q-Q plot
        from scipy import stats
        stats.probplot(residuals, dist="norm", plot=axes[1, 0])
        axes[1, 0].set_title('Q-Q Plot of Residuals')
        axes[1, 0].grid(True, alpha=0.3)
        
        # Actual vs Predicted
        axes[1, 1].scatter(y_true, y_pred, alpha=0.6)
        min_val = min(y_true.min(), y_pred.min())
        max_val = max(y_true.max(), y_pred.max())
        axes[1, 1].plot([min_val, max_val], [min_val, max_val], 'r--')
        axes[1, 1].set_xlabel('Actual Values')
        axes[1, 1].set_ylabel('Predicted Values')
        axes[1, 1].set_title('Actual vs Predicted')
        axes[1, 1].grid(True, alpha=0.3)
        
        plt.tight_layout()
        if save_path:
            plt.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
        plt.show()
    
    def plot_model_interpretability(self, model, X: np.ndarray, feature_names: List[str],
                                   method: str = 'shap', save_path: Optional[str] = None) -> None:
        """
        Plot model interpretability using SHAP or LIME
        
        Args:
            model: Trained model
            X: Input data
            feature_names: Names of features
            method: Interpretability method ('shap' or 'lime')
            save_path: Path to save the visualization
        """
        try:
            if method.lower() == 'shap':
                self._plot_shap_values(model, X, feature_names, save_path)
            elif method.lower() == 'lime':
                self._plot_lime_explanation(model, X, feature_names, save_path)
            else:
                print(f"Unsupported interpretability method: {method}")
        except ImportError:
            print(f"{method.upper()} is not installed. Install with: pip install {method}")
    
    def _plot_shap_values(self, model, X: np.ndarray, feature_names: List[str],
                         save_path: Optional[str] = None):
        """Plot SHAP values"""
        try:
            import shap
            
            # Create SHAP explainer
            explainer = shap.Explainer(model, X)
            shap_values = explainer(X)
            
            # Summary plot
            plt.figure(figsize=self.config.figure_size)
            shap.summary_plot(shap_values, X, feature_names=feature_names, show=False)
            plt.title('SHAP Feature Importance')
            
            if save_path:
                plt.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
            plt.show()
            
        except Exception as e:
            print(f"Error plotting SHAP values: {e}")
    
    def _plot_lime_explanation(self, model, X: np.ndarray, feature_names: List[str],
                              save_path: Optional[str] = None):
        """Plot LIME explanation"""
        try:
            import lime
            import lime.lime_tabular
            
            # Create LIME explainer
            explainer = lime.lime_tabular.LimeTabularExplainer(
                X, feature_names=feature_names, mode='regression'
            )
            
            # Explain a single instance
            idx = 0  # Explain first instance
            exp = explainer.explain_instance(X[idx], model.predict, num_features=10)
            
            # Plot explanation
            plt.figure(figsize=self.config.figure_size)
            exp.as_pyplot_figure()
            plt.title(f'LIME Explanation for Instance {idx}')
            
            if save_path:
                plt.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
            plt.show()
            
        except Exception as e:
            print(f"Error plotting LIME explanation: {e}")
    
    def create_interactive_dashboard(self, training_history: List[Dict[str, Any]],
                                   model_metrics: Dict[str, Any],
                                   save_path: Optional[str] = None) -> None:
        """
        Create an interactive dashboard with multiple visualizations
        
        Args:
            training_history: Training history data
            model_metrics: Model performance metrics
            save_path: Path to save the dashboard
        """
        if not PLOTLY_AVAILABLE:
            print("Plotly is required for interactive dashboard")
            return
        
        # Create subplots
        fig = make_subplots(
            rows=2, cols=2,
            subplot_titles=('Training Loss', 'Learning Rate', 'Metrics', 'Model Summary'),
            specs=[[{"secondary_y": False}, {"secondary_y": False}],
                   [{"secondary_y": False}, {"type": "table"}]]
        )
        
        epochs = list(range(1, len(training_history) + 1))
        
        # Training loss
        train_losses = [h.get('train_loss') for h in training_history]
        val_losses = [h.get('val_loss') for h in training_history]
        
        fig.add_trace(
            go.Scatter(x=epochs, y=train_losses, name='Training Loss', line=dict(color='blue')),
            row=1, col=1
        )
        
        if val_losses and all(v is not None for v in val_losses):
            fig.add_trace(
                go.Scatter(x=epochs, y=val_losses, name='Validation Loss', line=dict(color='red')),
                row=1, col=1
            )
        
        # Learning rate
        lrs = [h.get('learning_rate') for h in training_history if h.get('learning_rate')]
        if lrs:
            lr_epochs = epochs[:len(lrs)]
            fig.add_trace(
                go.Scatter(x=lr_epochs, y=lrs, name='Learning Rate', line=dict(color='green')),
                row=1, col=2
            )
        
        # Metrics
        if training_history and training_history[0].get('train_metrics'):
            for metric_name in list(training_history[0]['train_metrics'].keys())[:3]:  # First 3 metrics
                metric_values = [h['train_metrics'].get(metric_name) for h in training_history 
                               if h.get('train_metrics') and metric_name in h['train_metrics']]
                if metric_values and all(v is not None for v in metric_values):
                    fig.add_trace(
                        go.Scatter(x=epochs[:len(metric_values)], y=metric_values, 
                                 name=f'Training {metric_name}'),
                        row=2, col=1
                    )
        
        # Model summary table
        if model_metrics:
            fig.add_trace(
                go.Table(
                    header=dict(values=['Metric', 'Value']),
                    cells=dict(values=[list(model_metrics.keys()), list(model_metrics.values())])
                ),
                row=2, col=2
            )
        
        fig.update_layout(
            height=800,
            title_text="Model Training Dashboard",
            showlegend=True
        )
        
        if save_path:
            fig.write_html(save_path.replace('.png', '.html'))
        
        fig.show()
