"""
ML Visualizer

Comprehensive machine learning visualization toolkit with support for
data exploration, model analysis, and result visualization.
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import plotly.offline as pyo
from typing import Dict, List, Any, Optional, Tuple, Union
import warnings
from sklearn.metrics import confusion_matrix, classification_report
from sklearn.manifold import TSNE
from sklearn.decomposition import PCA
from sklearn.tree import plot_tree
import networkx as nx


class MLVisualizer:
    """
    Comprehensive ML visualization toolkit supporting various plot types
    for data exploration, model analysis, and results presentation.
    """
    
    def __init__(self, style: str = 'seaborn-v0_8', figsize: Tuple[int, int] = (12, 8)):
        """
        Initialize visualizer
        
        Args:
            style: Matplotlib style
            figsize: Default figure size
        """
        self.style = style
        self.figsize = figsize
        self.setup_style()
    
    def setup_style(self):
        """Setup visualization style"""
        try:
            plt.style.use(self.style)
        except:
            plt.style.use('default')
        
        sns.set_palette("husl")
        plt.rcParams['figure.figsize'] = self.figsize
    
    def plot_data_overview(self, df: pd.DataFrame, save_path: Optional[str] = None) -> None:
        """
        Create comprehensive data overview visualization
        
        Args:
            df: Input DataFrame
            save_path: Path to save the plot
        """
        fig, axes = plt.subplots(2, 3, figsize=(18, 12))
        fig.suptitle('Data Overview', fontsize=16, fontweight='bold')
        
        # Data types distribution
        dtype_counts = df.dtypes.value_counts()
        axes[0, 0].pie(dtype_counts.values, labels=dtype_counts.index, autopct='%1.1f%%')
        axes[0, 0].set_title('Data Types Distribution')
        
        # Missing values
        missing_counts = df.isnull().sum()
        missing_counts = missing_counts[missing_counts > 0].sort_values(ascending=False)
        if len(missing_counts) > 0:
            axes[0, 1].bar(range(len(missing_counts)), missing_counts.values)
            axes[0, 1].set_title('Missing Values')
            axes[0, 1].set_xlabel('Features')
            axes[0, 1].set_ylabel('Missing Count')
            axes[0, 1].set_xticks(range(len(missing_counts)))
            axes[0, 1].set_xticklabels(missing_counts.index, rotation=45)
        else:
            axes[0, 1].text(0.5, 0.5, 'No Missing Values', ha='center', va='center', 
                           transform=axes[0, 1].transAxes)
            axes[0, 1].set_title('Missing Values')
        
        # Numerical features distribution
        numerical_cols = df.select_dtypes(include=[np.number]).columns
        if len(numerical_cols) > 0:
            df[numerical_cols].hist(bins=20, ax=axes[0, 2], alpha=0.7)
            axes[0, 2].set_title('Numerical Features Distribution')
        else:
            axes[0, 2].text(0.5, 0.5, 'No Numerical Features', ha='center', va='center',
                           transform=axes[0, 2].transAxes)
            axes[0, 2].set_title('Numerical Features Distribution')
        
        # Correlation heatmap
        if len(numerical_cols) > 1:
            corr_matrix = df[numerical_cols].corr()
            sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', center=0, 
                       ax=axes[1, 0], fmt='.2f')
            axes[1, 0].set_title('Correlation Matrix')
        else:
            axes[1, 0].text(0.5, 0.5, 'Insufficient Features for Correlation', 
                           ha='center', va='center', transform=axes[1, 0].transAxes)
            axes[1, 0].set_title('Correlation Matrix')
        
        # Categorical features
        categorical_cols = df.select_dtypes(include=['object', 'category']).columns
        if len(categorical_cols) > 0:
            cat_counts = [df[col].nunique() for col in categorical_cols]
            axes[1, 1].bar(range(len(categorical_cols)), cat_counts)
            axes[1, 1].set_title('Categorical Features Cardinality')
            axes[1, 1].set_xlabel('Features')
            axes[1, 1].set_ylabel('Unique Values')
            axes[1, 1].set_xticks(range(len(categorical_cols)))
            axes[1, 1].set_xticklabels(categorical_cols, rotation=45)
        else:
            axes[1, 1].text(0.5, 0.5, 'No Categorical Features', ha='center', va='center',
                           transform=axes[1, 1].transAxes)
            axes[1, 1].set_title('Categorical Features Cardinality')
        
        # Data quality score
        quality_score = self._calculate_data_quality_score(df)
        axes[1, 2].bar(['Data Quality'], [quality_score])
        axes[1, 2].set_ylim(0, 100)
        axes[1, 2].set_title(f'Data Quality Score: {quality_score:.1f}')
        axes[1, 2].set_ylabel('Score')
        
        # Color code quality score
        if quality_score >= 80:
            color = 'green'
        elif quality_score >= 60:
            color = 'orange'
        else:
            color = 'red'
        axes[1, 2].patches[0].set_facecolor(color)
        
        plt.tight_layout()
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.show()
    
    def _calculate_data_quality_score(self, df: pd.DataFrame) -> float:
        """Calculate data quality score"""
        score = 100
        
        # Missing values penalty
        missing_percentage = (df.isnull().sum().sum() / (df.shape[0] * df.shape[1])) * 100
        score -= missing_percentage
        
        # Duplicate rows penalty
        duplicate_percentage = (df.duplicated().sum() / len(df)) * 100
        score -= duplicate_percentage * 2
        
        # Constant columns penalty
        constant_cols = [col for col in df.columns if df[col].nunique() == 1]
        score -= len(constant_cols) * 5
        
        return max(0, score)
    
    def plot_feature_importance(self, importance_dict: Dict[str, float], 
                               top_n: int = 20, save_path: Optional[str] = None) -> None:
        """
        Plot feature importance
        
        Args:
            importance_dict: Dictionary of feature names and importance scores
            top_n: Number of top features to display
            save_path: Path to save the plot
        """
        # Sort features by importance
        sorted_features = sorted(importance_dict.items(), key=lambda x: x[1], reverse=True)
        top_features = sorted_features[:top_n]
        
        features, importance = zip(*top_features)
        
        plt.figure(figsize=self.figsize)
        bars = plt.barh(range(len(features)), importance, color='skyblue', edgecolor='navy')
        
        # Add value labels
        for i, (bar, score) in enumerate(zip(bars, importance)):
            plt.text(bar.get_width() + 0.001, bar.get_y() + bar.get_height()/2,
                    f'{score:.3f}', ha='left', va='center', fontsize=10)
        
        plt.yticks(range(len(features)), features)
        plt.xlabel('Importance Score')
        plt.title(f'Top {top_n} Feature Importance')
        plt.gca().invert_yaxis()
        plt.grid(axis='x', alpha=0.3)
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.show()
    
    def plot_model_comparison(self, model_results: Dict[str, Dict[str, float]], 
                            metric: str = 'accuracy', save_path: Optional[str] = None) -> None:
        """
        Compare multiple models performance
        
        Args:
            model_results: Dictionary of model results
            metric: Metric to compare
            save_path: Path to save the plot
        """
        models = list(model_results.keys())
        scores = [model_results[model].get(metric, 0) for model in models]
        
        plt.figure(figsize=self.figsize)
        bars = plt.bar(models, scores, color='lightcoral', edgecolor='darkred')
        
        # Add value labels
        for bar, score in zip(bars, scores):
            plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01,
                    f'{score:.3f}', ha='center', va='bottom', fontsize=12)
        
        plt.xlabel('Models')
        plt.ylabel(metric.replace('_', ' ').title())
        plt.title(f'Model Comparison - {metric.replace("_", " ").title()}')
        plt.xticks(rotation=45)
        plt.grid(axis='y', alpha=0.3)
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, 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 plot
        """
        cm = confusion_matrix(y_true, y_pred)
        
        plt.figure(figsize=self.figsize)
        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=300, bbox_inches='tight')
        plt.show()
    
    def plot_learning_curves(self, history: Dict[str, List[float]], 
                           save_path: Optional[str] = None) -> None:
        """
        Plot learning curves from training history
        
        Args:
            history: Training history dictionary
            save_path: Path to save the plot
        """
        epochs = list(range(1, len(history.get('loss', [])) + 1))
        
        fig, axes = plt.subplots(2, 2, figsize=(15, 10))
        
        # Loss curves
        if 'loss' in history:
            axes[0, 0].plot(epochs, history['loss'], label='Training Loss', linewidth=2)
        if 'val_loss' in history:
            axes[0, 0].plot(epochs, history['val_loss'], label='Validation Loss', linewidth=2)
        axes[0, 0].set_title('Loss Curves')
        axes[0, 0].set_xlabel('Epoch')
        axes[0, 0].set_ylabel('Loss')
        axes[0, 0].legend()
        axes[0, 0].grid(True, alpha=0.3)
        
        # Accuracy curves
        if 'accuracy' in history:
            axes[0, 1].plot(epochs, history['accuracy'], label='Training Accuracy', linewidth=2)
        if 'val_accuracy' in history:
            axes[0, 1].plot(epochs, history['val_accuracy'], label='Validation Accuracy', linewidth=2)
        axes[0, 1].set_title('Accuracy Curves')
        axes[0, 1].set_xlabel('Epoch')
        axes[0, 1].set_ylabel('Accuracy')
        axes[0, 1].legend()
        axes[0, 1].grid(True, alpha=0.3)
        
        # Learning rate (if available)
        if 'lr' in history:
            axes[1, 0].plot(epochs, history['lr'], label='Learning Rate', linewidth=2, color='orange')
            axes[1, 0].set_title('Learning Rate Schedule')
            axes[1, 0].set_xlabel('Epoch')
            axes[1, 0].set_ylabel('Learning Rate')
            axes[1, 0].legend()
            axes[1, 0].grid(True, alpha=0.3)
        
        # Metrics comparison
        metrics = [k for k in history.keys() if k not in ['loss', 'val_loss', 'lr', 'accuracy', 'val_accuracy']]
        if metrics:
            for metric in metrics[:3]:  # Plot first 3 additional metrics
                if metric in history:
                    axes[1, 1].plot(epochs, history[metric], label=metric.replace('_', ' ').title(), linewidth=2)
            axes[1, 1].set_title('Additional Metrics')
            axes[1, 1].set_xlabel('Epoch')
            axes[1, 1].set_ylabel('Metric Value')
            axes[1, 1].legend()
            axes[1, 1].grid(True, alpha=0.3)
        
        plt.tight_layout()
        if save_path:
            plt.savefig(save_path, dpi=300, 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 plot
        """
        residuals = y_true - y_pred
        
        fig, axes = plt.subplots(2, 2, figsize=(15, 10))
        
        # Residuals vs Predicted
        axes[0, 0].scatter(y_pred, residuals, alpha=0.6, color='blue')
        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, color='green')
        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=300, bbox_inches='tight')
        plt.show()
    
    def plot_feature_distribution(self, df: pd.DataFrame, features: List[str],
                                 target: Optional[str] = None,
                                 save_path: Optional[str] = None) -> None:
        """
        Plot distribution of features with optional target separation
        
        Args:
            df: Input DataFrame
            features: List of features to plot
            target: Target column for color coding
            save_path: Path to save the plot
        """
        n_features = len(features)
        n_cols = min(3, n_features)
        n_rows = (n_features + n_cols - 1) // n_cols
        
        fig, axes = plt.subplots(n_rows, n_cols, figsize=(n_cols * 5, n_rows * 4))
        if n_features == 1:
            axes = [axes]
        elif n_rows == 1:
            axes = axes.reshape(1, -1)
        
        for i, feature in enumerate(features):
            row = i // n_cols
            col = i % n_cols
            ax = axes[row, col] if n_rows > 1 else axes[col]
            
            if target and target in df.columns:
                # Plot with target separation
                for target_val in df[target].unique():
                    subset = df[df[target] == target_val][feature]
                    ax.hist(subset, alpha=0.6, label=str(target_val), bins=20)
                ax.legend()
            else:
                # Plot single distribution
                ax.hist(df[feature], alpha=0.7, bins=20, color='skyblue', edgecolor='black')
            
            ax.set_title(f'Distribution of {feature}')
            ax.set_xlabel(feature)
            ax.set_ylabel('Frequency')
            ax.grid(True, alpha=0.3)
        
        # Hide empty subplots
        for i in range(n_features, n_rows * n_cols):
            row = i // n_cols
            col = i % n_cols
            if n_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=300, bbox_inches='tight')
        plt.show()
    
    def plot_correlation_matrix(self, df: pd.DataFrame, 
                               save_path: Optional[str] = None) -> None:
        """
        Plot correlation matrix heatmap
        
        Args:
            df: Input DataFrame
            save_path: Path to save the plot
        """
        numerical_cols = df.select_dtypes(include=[np.number]).columns
        
        if len(numerical_cols) < 2:
            print("Need at least 2 numerical features for correlation matrix")
            return
        
        corr_matrix = df[numerical_cols].corr()
        
        plt.figure(figsize=(max(10, len(numerical_cols)), max(8, len(numerical_cols))))
        sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', center=0,
                   square=True, fmt='.2f', cbar_kws={"shrink": .8})
        plt.title('Feature Correlation Matrix')
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.show()
    
    def plot_dimensionality_reduction(self, X: np.ndarray, y: Optional[np.ndarray] = None,
                                     method: str = 'pca', n_components: int = 2,
                                     save_path: Optional[str] = None) -> None:
        """
        Plot dimensionality reduction visualization
        
        Args:
            X: Input features
            y: Target labels (optional)
            method: Reduction method ('pca', 'tsne')
            n_components: Number of components
            save_path: Path to save the plot
        """
        if method.lower() == 'pca':
            reducer = PCA(n_components=n_components, random_state=42)
            title = 'PCA Visualization'
        elif method.lower() == 'tsne':
            reducer = TSNE(n_components=n_components, random_state=42)
            title = 't-SNE Visualization'
        else:
            raise ValueError("Method must be 'pca' or 'tsne'")
        
        X_reduced = reducer.fit_transform(X)
        
        if n_components == 2:
            plt.figure(figsize=self.figsize)
            
            if y is not None:
                scatter = plt.scatter(X_reduced[:, 0], X_reduced[:, 1], c=y, 
                                   cmap='viridis', alpha=0.7)
                plt.colorbar(scatter, label='Target')
            else:
                plt.scatter(X_reduced[:, 0], X_reduced[:, 1], alpha=0.7)
            
            plt.xlabel(f'{method.upper()} Component 1')
            plt.ylabel(f'{method.upper()} Component 2')
            plt.title(title)
            plt.grid(True, alpha=0.3)
            
        elif n_components == 3:
            fig = plt.figure(figsize=self.figsize)
            ax = fig.add_subplot(111, projection='3d')
            
            if y is not None:
                scatter = ax.scatter(X_reduced[:, 0], X_reduced[:, 1], X_reduced[:, 2], 
                                   c=y, cmap='viridis', alpha=0.7)
                plt.colorbar(scatter, label='Target')
            else:
                ax.scatter(X_reduced[:, 0], X_reduced[:, 1], X_reduced[:, 2], alpha=0.7)
            
            ax.set_xlabel(f'{method.upper()} Component 1')
            ax.set_ylabel(f'{method.upper()} Component 2')
            ax.set_zlabel(f'{method.upper()} Component 3')
            ax.set_title(title)
        
        plt.tight_layout()
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.show()
    
    def plot_interactive_scatter(self, df: pd.DataFrame, x: str, y: str,
                               color: Optional[str] = None, size: Optional[str] = None,
                               save_path: Optional[str] = None) -> None:
        """
        Create interactive scatter plot using Plotly
        
        Args:
            df: Input DataFrame
            x: X-axis column
            y: Y-axis column
            color: Color column (optional)
            size: Size column (optional)
            save_path: Path to save the plot
        """
        fig = px.scatter(df, x=x, y=y, color=color, size=size,
                        hover_data=[col for col in df.columns if col not in [x, y]],
                        title=f'Interactive Scatter: {x} vs {y}')
        
        fig.update_layout(
            width=800,
            height=600,
            showlegend=True
        )
        
        if save_path:
            if save_path.endswith('.html'):
                fig.write_html(save_path)
            else:
                fig.write_image(save_path)
        
        fig.show()
    
    def plot_interactive_histogram(self, df: pd.DataFrame, column: str,
                                  bins: int = 30, save_path: Optional[str] = None) -> None:
        """
        Create interactive histogram using Plotly
        
        Args:
            df: Input DataFrame
            column: Column to plot
            bins: Number of bins
            save_path: Path to save the plot
        """
        fig = px.histogram(df, x=column, nbins=bins,
                          title=f'Distribution of {column}')
        
        fig.update_layout(
            width=800,
            height=600,
            showlegend=False
        )
        
        if save_path:
            if save_path.endswith('.html'):
                fig.write_html(save_path)
            else:
                fig.write_image(save_path)
        
        fig.show()
    
    def plot_model_architecture(self, layers: List[Dict[str, Any]],
                               save_path: Optional[str] = None) -> None:
        """
        Visualize neural network architecture
        
        Args:
            layers: List of layer configurations
            save_path: Path to save the plot
        """
        fig, ax = plt.subplots(figsize=(15, 8))
        
        # Calculate positions
        max_units = max([layer.get('units', layer.get('filters', 10)) for layer in layers])
        layer_spacing = 3.0
        
        for i, layer in enumerate(layers):
            x = i * layer_spacing
            layer_type = layer.get('type', 'dense').lower()
            
            if layer_type in ['dense', 'Dense']:
                units = layer.get('units', 10)
                # Draw neurons
                for j in range(min(units, 10)):  # Limit to 10 for visualization
                    y = (j - units/2) * 0.3
                    circle = plt.Circle((x, y), 0.1, color='lightblue', ec='darkblue')
                    ax.add_patch(circle)
                
                # Draw layer box
                rect = plt.Rectangle((x - 0.5, -units * 0.15), 1.0, units * 0.3,
                                   fill=False, ec='blue', linewidth=2)
                ax.add_patch(rect)
                
            elif layer_type in ['conv2d', 'Conv2D']:
                rect = plt.Rectangle((x - 0.5, -1), 1.0, 2.0,
                                   fill=True, facecolor='lightgreen', ec='darkgreen')
                ax.add_patch(rect)
                ax.text(x, 0, f"Conv\n{layer.get('filters', 32)}", 
                       ha='center', va='center', fontsize=10)
                
            elif layer_type in ['dropout', 'Dropout']:
                rect = plt.Rectangle((x - 0.5, -1), 1.0, 2.0,
                                   fill=True, facecolor='lightyellow', ec='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=10)
                
            elif layer_type in ['lstm', 'LSTM']:
                rect = plt.Rectangle((x - 0.5, -1), 1.0, 2.0,
                                   fill=True, facecolor='lightcoral', ec='red')
                ax.add_patch(rect)
                ax.text(x, 0, f"LSTM\n{layer.get('units', 64)}", 
                       ha='center', va='center', fontsize=10)
            
            # Add layer label
            ax.text(x, -2, layer.get('name', f'Layer {i}'), 
                   ha='center', va='top', fontsize=10, fontweight='bold')
        
        # Draw connections
        for i in range(len(layers) - 1):
            x1 = i * layer_spacing
            x2 = (i + 1) * layer_spacing
            ax.plot([x1 + 0.5, x2 - 0.5], [0, 0], 'gray', alpha=0.5, linewidth=1)
        
        ax.set_xlim(-1, len(layers) * layer_spacing)
        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=300, bbox_inches='tight')
        plt.show()
    
    def create_summary_report(self, df: pd.DataFrame, model_results: Optional[Dict] = None,
                           save_path: Optional[str] = None) -> Dict[str, Any]:
        """
        Create comprehensive summary report
        
        Args:
            df: Input DataFrame
            model_results: Model performance results
            save_path: Path to save the report
            
        Returns:
            Summary report dictionary
        """
        report = {
            'data_summary': {
                'shape': df.shape,
                'columns': list(df.columns),
                'dtypes': df.dtypes.to_dict(),
                'missing_values': df.isnull().sum().to_dict(),
                'data_quality_score': self._calculate_data_quality_score(df)
            },
            'feature_analysis': {
                'numerical_features': df.select_dtypes(include=[np.number]).columns.tolist(),
                'categorical_features': df.select_dtypes(include=['object', 'category']).columns.tolist(),
                'correlation_matrix': df.select_dtypes(include=[np.number]).corr().to_dict()
            }
        }
        
        if model_results:
            report['model_results'] = model_results
        
        if save_path:
            with open(save_path, 'w') as f:
                json.dump(report, f, indent=2, default=str)
        
        return report
