"""
Model Evaluator

Comprehensive model evaluation with advanced metrics, visualization,
and AI-powered performance analysis.
"""

import pandas as pd
import numpy as np
from typing import Dict, List, Any, Optional, Tuple
from sklearn.metrics import (
    accuracy_score, precision_score, recall_score, f1_score, roc_auc_score,
    confusion_matrix, classification_report, roc_curve, precision_recall_curve,
    mean_squared_error, mean_absolute_error, r2_score, explained_variance_score,
    silhouette_score, calinski_harabasz_score, davies_bouldin_score
)
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import stats
import warnings


class ModelEvaluator:
    """
    Comprehensive model evaluator with advanced metrics and AI-powered
    performance analysis and recommendations.
    """
    
    def __init__(self):
        self.evaluation_results = {}
        self.comparison_results = {}
        
    def evaluate_classification(self, y_true: np.ndarray, y_pred: np.ndarray,
                              y_pred_proba: Optional[np.ndarray] = None,
                              class_names: Optional[List[str]] = None) -> Dict[str, Any]:
        """
        Comprehensive classification model evaluation
        
        Args:
            y_true: True labels
            y_pred: Predicted labels
            y_pred_proba: Predicted probabilities (optional)
            class_names: Names of classes (optional)
            
        Returns:
            Dictionary containing comprehensive evaluation metrics
        """
        results = {}
        
        # Basic metrics
        results['accuracy'] = accuracy_score(y_true, y_pred)
        results['precision_macro'] = precision_score(y_true, y_pred, average='macro', zero_division=0)
        results['precision_micro'] = precision_score(y_true, y_pred, average='micro', zero_division=0)
        results['precision_weighted'] = precision_score(y_true, y_pred, average='weighted', zero_division=0)
        
        results['recall_macro'] = recall_score(y_true, y_pred, average='macro', zero_division=0)
        results['recall_micro'] = recall_score(y_true, y_pred, average='micro', zero_division=0)
        results['recall_weighted'] = recall_score(y_true, y_pred, average='weighted', zero_division=0)
        
        results['f1_macro'] = f1_score(y_true, y_pred, average='macro', zero_division=0)
        results['f1_micro'] = f1_score(y_true, y_pred, average='micro', zero_division=0)
        results['f1_weighted'] = f1_score(y_true, y_pred, average='weighted', zero_division=0)
        
        # Per-class metrics
        results['classification_report'] = classification_report(
            y_true, y_pred, target_names=class_names, output_dict=True, zero_division=0
        )
        
        # Confusion matrix
        cm = confusion_matrix(y_true, y_pred)
        results['confusion_matrix'] = cm.tolist()
        
        # ROC AUC (if probabilities provided)
        if y_pred_proba is not None:
            try:
                if len(np.unique(y_true)) == 2:  # Binary classification
                    results['roc_auc'] = roc_auc_score(y_true, y_pred_proba[:, 1])
                    
                    # ROC curve data
                    fpr, tpr, _ = roc_curve(y_true, y_pred_proba[:, 1])
                    results['roc_curve'] = {'fpr': fpr.tolist(), 'tpr': tpr.tolist()}
                    
                    # Precision-Recall curve data
                    precision, recall, _ = precision_recall_curve(y_true, y_pred_proba[:, 1])
                    results['pr_curve'] = {'precision': precision.tolist(), 'recall': recall.tolist()}
                    
                else:  # Multi-class classification
                    results['roc_auc_ovr'] = roc_auc_score(y_true, y_pred_proba, multi_class='ovr', average='macro')
                    results['roc_auc_ovo'] = roc_auc_score(y_true, y_pred_proba, multi_class='ovo', average='macro')
                    
            except Exception as e:
                results['roc_auc_error'] = str(e)
        
        # Class balance analysis
        unique, counts = np.unique(y_true, return_counts=True)
        class_distribution = dict(zip(unique, counts))
        results['class_distribution'] = class_distribution
        
        # Performance analysis
        results['performance_analysis'] = self._analyze_classification_performance(results)
        
        return results
    
    def evaluate_regression(self, y_true: np.ndarray, y_pred: np.ndarray) -> Dict[str, Any]:
        """
        Comprehensive regression model evaluation
        
        Args:
            y_true: True values
            y_pred: Predicted values
            
        Returns:
            Dictionary containing comprehensive evaluation metrics
        """
        results = {}
        
        # Basic metrics
        results['mse'] = mean_squared_error(y_true, y_pred)
        results['rmse'] = np.sqrt(results['mse'])
        results['mae'] = mean_absolute_error(y_true, y_pred)
        results['r2'] = r2_score(y_true, y_pred)
        results['explained_variance'] = explained_variance_score(y_true, y_pred)
        
        # Additional metrics
        results['mean_absolute_percentage_error'] = np.mean(np.abs((y_true - y_pred) / y_true)) * 100
        results['max_error'] = np.max(np.abs(y_true - y_pred))
        
        # Residual analysis
        residuals = y_true - y_pred
        results['residuals'] = residuals.tolist()
        results['residual_mean'] = np.mean(residuals)
        results['residual_std'] = np.std(residuals)
        
        # Normality test for residuals
        _, p_value = stats.normaltest(residuals)
        results['residual_normality_p_value'] = p_value
        results['residuals_are_normal'] = p_value > 0.05
        
        # Heteroscedasticity check (simple version)
        results['heteroscedasticity'] = self._check_heteroscedasticity(y_pred, residuals)
        
        # Performance analysis
        results['performance_analysis'] = self._analyze_regression_performance(results)
        
        return results
    
    def evaluate_clustering(self, X: np.ndarray, labels: np.ndarray, 
                           true_labels: Optional[np.ndarray] = None) -> Dict[str, Any]:
        """
        Comprehensive clustering model evaluation
        
        Args:
            X: Feature matrix
            labels: Cluster labels
            true_labels: True labels (optional, for supervised evaluation)
            
        Returns:
            Dictionary containing comprehensive evaluation metrics
        """
        results = {}
        
        # Internal metrics (unsupervised)
        try:
            results['silhouette_score'] = silhouette_score(X, labels)
        except:
            results['silhouette_score'] = None
        
        try:
            results['calinski_harabasz_score'] = calinski_harabasz_score(X, labels)
        except:
            results['calinski_harabasz_score'] = None
        
        try:
            results['davies_bouldin_score'] = davies_bouldin_score(X, labels)
        except:
            results['davies_bouldin_score'] = None
        
        # Cluster analysis
        unique_labels, counts = np.unique(labels, return_counts=True)
        results['cluster_sizes'] = dict(zip(unique_labels, counts))
        results['n_clusters'] = len(unique_labels)
        results['cluster_balance'] = np.std(counts) / np.mean(counts) if np.mean(counts) > 0 else 0
        
        # External metrics (if true labels available)
        if true_labels is not None:
            try:
                from sklearn.metrics import adjusted_rand_score, normalized_mutual_info_score
                results['adjusted_rand_score'] = adjusted_rand_score(true_labels, labels)
                results['normalized_mutual_info'] = normalized_mutual_info_score(true_labels, labels)
            except:
                pass
        
        # Performance analysis
        results['performance_analysis'] = self._analyze_clustering_performance(results)
        
        return results
    
    def compare_models(self, model_results: Dict[str, Dict[str, Any]], 
                      problem_type: str) -> Dict[str, Any]:
        """
        Compare multiple models and provide rankings
        
        Args:
            model_results: Dictionary of model evaluation results
            problem_type: 'classification', 'regression', or 'clustering'
            
        Returns:
            Dictionary containing model comparisons and rankings
        """
        comparison = {
            'model_rankings': [],
            'metric_comparison': {},
            'statistical_tests': {},
            'recommendations': []
        }
        
        # Extract key metrics for comparison
        if problem_type == 'classification':
            key_metrics = ['accuracy', 'f1_weighted', 'precision_weighted', 'recall_weighted']
            if any('roc_auc' in results for results in model_results.values()):
                key_metrics.append('roc_auc')
        elif problem_type == 'regression':
            key_metrics = ['r2', 'rmse', 'mae', 'explained_variance']
        else:  # clustering
            key_metrics = ['silhouette_score', 'calinski_harabasz_score']
        
        # Create comparison table
        comparison_table = {}
        for model_name, results in model_results.items():
            comparison_table[model_name] = {}
            for metric in key_metrics:
                if metric in results:
                    comparison_table[model_name][metric] = results[metric]
        
        comparison['metric_comparison'] = comparison_table
        
        # Rank models
        rankings = []
        for model_name in model_results.keys():
            score = 0
            count = 0
            
            for metric in key_metrics:
                if metric in model_results[model_name]:
                    if metric in ['rmse', 'mae', 'davies_bouldin_score']:
                        # Lower is better
                        score -= model_results[model_name][metric]
                    else:
                        # Higher is better
                        score += model_results[model_name][metric]
                    count += 1
            
            if count > 0:
                rankings.append({
                    'model': model_name,
                    'composite_score': score / count,
                    'metrics': comparison_table[model_name]
                })
        
        # Sort by composite score
        if problem_type in ['classification', 'clustering']:
            rankings.sort(key=lambda x: x['composite_score'], reverse=True)
        else:  # regression
            rankings.sort(key=lambda x: x['composite_score'], reverse=True)
        
        comparison['model_rankings'] = rankings
        
        # Statistical significance tests (if we have predictions)
        comparison['statistical_tests'] = self._perform_statistical_tests(model_results, problem_type)
        
        # Generate recommendations
        comparison['recommendations'] = self._generate_comparison_recommendations(rankings, problem_type)
        
        self.comparison_results = comparison
        return comparison
    
    def _analyze_classification_performance(self, results: Dict[str, Any]) -> Dict[str, Any]:
        """Analyze classification performance and provide insights"""
        analysis = {}
        
        accuracy = results.get('accuracy', 0)
        f1_weighted = results.get('f1_weighted', 0)
        
        # Performance tier
        if accuracy >= 0.9:
            analysis['performance_tier'] = 'Excellent'
        elif accuracy >= 0.8:
            analysis['performance_tier'] = 'Good'
        elif accuracy >= 0.7:
            analysis['performance_tier'] = 'Fair'
        else:
            analysis['performance_tier'] = 'Poor'
        
        # Class imbalance issues
        class_dist = results.get('class_distribution', {})
        if class_dist:
            max_count = max(class_dist.values())
            min_count = min(class_dist.values())
            imbalance_ratio = max_count / min_count if min_count > 0 else float('inf')
            
            if imbalance_ratio > 10:
                analysis['class_imbalance'] = 'Severe'
            elif imbalance_ratio > 3:
                analysis['class_imbalance'] = 'Moderate'
            else:
                analysis['class_imbalance'] = 'Balanced'
        
        # Potential issues
        issues = []
        if accuracy < 0.6:
            issues.append('Low accuracy suggests model may be underperforming')
        
        if results.get('precision_macro', 0) < results.get('precision_weighted', 0) * 0.8:
            issues.append('Macro precision significantly lower than weighted precision - possible class imbalance issues')
        
        if results.get('recall_macro', 0) < results.get('recall_weighted', 0) * 0.8:
            issues.append('Macro recall significantly lower than weighted recall - possible class imbalance issues')
        
        analysis['potential_issues'] = issues
        
        return analysis
    
    def _analyze_regression_performance(self, results: Dict[str, Any]) -> Dict[str, Any]:
        """Analyze regression performance and provide insights"""
        analysis = {}
        
        r2 = results.get('r2', 0)
        rmse = results.get('rmse', 0)
        
        # Performance tier
        if r2 >= 0.9:
            analysis['performance_tier'] = 'Excellent'
        elif r2 >= 0.7:
            analysis['performance_tier'] = 'Good'
        elif r2 >= 0.5:
            analysis['performance_tier'] = 'Fair'
        else:
            analysis['performance_tier'] = 'Poor'
        
        # Residual analysis
        if results.get('residuals_are_normal', False):
            analysis['residual_normality'] = 'Normal'
        else:
            analysis['residual_normality'] = 'Non-normal'
        
        if results.get('heteroscedasticity', False):
            analysis['heteroscedasticity'] = 'Present'
        else:
            analysis['heteroscedasticity'] = 'Absent'
        
        # Potential issues
        issues = []
        if r2 < 0.3:
            issues.append('Low R² suggests poor model fit')
        
        if results.get('residual_mean', 0) != 0:
            issues.append('Non-zero residual mean indicates bias in predictions')
        
        if not results.get('residuals_are_normal', True):
            issues.append('Non-normal residuals suggest model assumptions violated')
        
        analysis['potential_issues'] = issues
        
        return analysis
    
    def _analyze_clustering_performance(self, results: Dict[str, Any]) -> Dict[str, Any]:
        """Analyze clustering performance and provide insights"""
        analysis = {}
        
        silhouette = results.get('silhouette_score', 0)
        n_clusters = results.get('n_clusters', 0)
        cluster_balance = results.get('cluster_balance', 0)
        
        # Performance tier based on silhouette score
        if silhouette >= 0.7:
            analysis['performance_tier'] = 'Excellent'
        elif silhouette >= 0.5:
            analysis['performance_tier'] = 'Good'
        elif silhouette >= 0.25:
            analysis['performance_tier'] = 'Fair'
        else:
            analysis['performance_tier'] = 'Poor'
        
        # Cluster analysis
        if cluster_balance > 1.5:
            analysis['cluster_balance'] = 'Imbalanced'
        else:
            analysis['cluster_balance'] = 'Balanced'
        
        # Potential issues
        issues = []
        if silhouette < 0.25:
            issues.append('Low silhouette score suggests poor clustering structure')
        
        if n_clusters == 1:
            issues.append('Only one cluster found - consider different parameters')
        
        if cluster_balance > 2.0:
            issues.append('Highly imbalanced clusters - consider different clustering approach')
        
        analysis['potential_issues'] = issues
        
        return analysis
    
    def _check_heteroscedasticity(self, y_pred: np.ndarray, residuals: np.ndarray) -> bool:
        """Simple check for heteroscedasticity"""
        # Correlate absolute residuals with predicted values
        correlation = np.corrcoef(np.abs(residuals), y_pred)[0, 1]
        return abs(correlation) > 0.3
    
    def _perform_statistical_tests(self, model_results: Dict[str, Any], 
                                 problem_type: str) -> Dict[str, Any]:
        """Perform statistical tests to compare models"""
        tests = {}
        
        # This is a simplified version - in practice, you'd need the actual predictions
        # to perform proper statistical tests like paired t-tests
        
        if len(model_results) >= 2:
            models = list(model_results.keys())
            
            # Simple comparison of key metrics
            if problem_type == 'classification':
                metric = 'accuracy'
            elif problem_type == 'regression':
                metric = 'r2'
            else:
                metric = 'silhouette_score'
            
            for i in range(len(models)):
                for j in range(i + 1, len(models)):
                    model1, model2 = models[i], models[j]
                    score1 = model_results[model1].get(metric, 0)
                    score2 = model_results[model2].get(metric, 0)
                    
                    difference = abs(score1 - score2)
                    tests[f"{model1}_vs_{model2}"] = {
                        'metric': metric,
                        'difference': difference,
                        'better_model': model1 if score1 > score2 else model2
                    }
        
        return tests
    
    def _generate_comparison_recommendations(self, rankings: List[Dict[str, Any]], 
                                          problem_type: str) -> List[str]:
        """Generate recommendations based on model comparison"""
        recommendations = []
        
        if not rankings:
            return ["No models to compare"]
        
        best_model = rankings[0]
        second_best = rankings[1] if len(rankings) > 1 else None
        
        recommendations.append(f"Best performing model: {best_model['model']}")
        
        # Check if the best model is significantly better
        if second_best:
            score_diff = best_model['composite_score'] - second_best['composite_score']
            if score_diff < 0.05:  # Small difference
                recommendations.append(f"Performance difference between top models is small ({score_diff:.3f})")
                recommendations.append("Consider ensemble methods or choose based on other factors")
            elif score_diff > 0.2:  # Large difference
                recommendations.append(f"{best_model['model']} significantly outperforms other models")
        
        # Problem-specific recommendations
        if problem_type == 'classification':
            if best_model['metrics'].get('accuracy', 0) < 0.8:
                recommendations.append("Consider feature engineering or ensemble methods to improve accuracy")
        elif problem_type == 'regression':
            if best_model['metrics'].get('r2', 0) < 0.7:
                recommendations.append("Consider non-linear models or feature transformation")
        
        return recommendations
    
    def plot_evaluation_results(self, results: Dict[str, Any], problem_type: str,
                               save_path: Optional[str] = None):
        """Plot evaluation results"""
        try:
            if problem_type == 'classification':
                self._plot_classification_results(results, save_path)
            elif problem_type == 'regression':
                self._plot_regression_results(results, save_path)
            else:
                self._plot_clustering_results(results, save_path)
        except Exception as e:
            print(f"Plotting failed: {e}")
    
    def _plot_classification_results(self, results: Dict[str, Any], save_path: Optional[str]):
        """Plot classification evaluation results"""
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        
        # Confusion matrix
        cm = np.array(results['confusion_matrix'])
        sns.heatmap(cm, annot=True, fmt='d', ax=axes[0, 0], cmap='Blues')
        axes[0, 0].set_title('Confusion Matrix')
        axes[0, 0].set_xlabel('Predicted')
        axes[0, 0].set_ylabel('Actual')
        
        # Metrics bar chart
        metrics = ['accuracy', 'precision_macro', 'recall_macro', 'f1_macro']
        values = [results.get(m, 0) for m in metrics]
        axes[0, 1].bar(metrics, values)
        axes[0, 1].set_title('Classification Metrics')
        axes[0, 1].set_ylabel('Score')
        axes[0, 1].tick_params(axis='x', rotation=45)
        
        # ROC curve (if available)
        if 'roc_curve' in results:
            fpr = results['roc_curve']['fpr']
            tpr = results['roc_curve']['tpr']
            axes[1, 0].plot(fpr, tpr, label=f'ROC AUC = {results.get("roc_auc", 0):.3f}')
            axes[1, 0].plot([0, 1], [0, 1], 'k--')
            axes[1, 0].set_xlabel('False Positive Rate')
            axes[1, 0].set_ylabel('True Positive Rate')
            axes[1, 0].set_title('ROC Curve')
            axes[1, 0].legend()
        
        # Class distribution
        if 'class_distribution' in results:
            class_dist = results['class_distribution']
            axes[1, 1].bar(class_dist.keys(), class_dist.values())
            axes[1, 1].set_title('Class Distribution')
            axes[1, 1].set_xlabel('Class')
            axes[1, 1].set_ylabel('Count')
        
        plt.tight_layout()
        if save_path:
            plt.savefig(save_path)
        plt.show()
    
    def _plot_regression_results(self, results: Dict[str, Any], save_path: Optional[str]):
        """Plot regression evaluation results"""
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        
        # Metrics bar chart
        metrics = ['r2', 'explained_variance']
        values = [results.get(m, 0) for m in metrics]
        axes[0, 0].bar(metrics, values)
        axes[0, 0].set_title('R² and Explained Variance')
        axes[0, 0].set_ylabel('Score')
        axes[0, 0].set_ylim(0, 1)
        
        # Error metrics
        error_metrics = ['rmse', 'mae']
        error_values = [results.get(m, 0) for m in error_metrics]
        axes[0, 1].bar(error_metrics, error_values)
        axes[0, 1].set_title('Error Metrics')
        axes[0, 1].set_ylabel('Error')
        
        # Residuals histogram
        residuals = results.get('residuals', [])
        if residuals:
            axes[1, 0].hist(residuals, bins=30, alpha=0.7)
            axes[1, 0].set_title('Residuals Distribution')
            axes[1, 0].set_xlabel('Residual')
            axes[1, 0].set_ylabel('Frequency')
        
        # Residuals vs predicted (if we had predicted values)
        # This would need actual predictions to plot properly
        axes[1, 1].text(0.5, 0.5, 'Residuals vs Predicted\n(Requires actual predictions)', 
                       ha='center', va='center', transform=axes[1, 1].transAxes)
        axes[1, 1].set_title('Residuals Analysis')
        
        plt.tight_layout()
        if save_path:
            plt.savefig(save_path)
        plt.show()
    
    def _plot_clustering_results(self, results: Dict[str, Any], save_path: Optional[str]):
        """Plot clustering evaluation results"""
        fig, axes = plt.subplots(1, 2, figsize=(12, 5))
        
        # Metrics bar chart
        metrics = ['silhouette_score', 'calinski_harabasz_score']
        # Normalize calinski_harabasz for better visualization
        ch_score = results.get('calinski_harabasz_score', 0)
        ch_normalized = ch_score / max(results.get('calinski_harabasz_score', 1), 1) * 0.5
        
        values = [results.get('silhouette_score', 0), ch_normalized]
        axes[0].bar(metrics, values)
        axes[0].set_title('Clustering Metrics')
        axes[0].set_ylabel('Score')
        axes[0].tick_params(axis='x', rotation=45)
        
        # Cluster sizes
        if 'cluster_sizes' in results:
            cluster_sizes = results['cluster_sizes']
            axes[1].bar(cluster_sizes.keys(), cluster_sizes.values())
            axes[1].set_title('Cluster Sizes')
            axes[1].set_xlabel('Cluster')
            axes[1].set_ylabel('Size')
        
        plt.tight_layout()
        if save_path:
            plt.savefig(save_path)
        plt.show()
