"""
Model Comparison and Evaluation Utilities

This module provides comprehensive comparison tools for evaluating AMBER
against baseline models and analyzing disambiguation performance.
"""

import numpy as np
from typing import List, Dict, Tuple, Union, Optional
from scipy.spatial.distance import cosine
from collections import defaultdict
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from .metrics import MetricsCalculator, ContextSensitivityScore


class AMBERComparator:
    """
    Comprehensive comparison and evaluation toolkit for AMBER model
    """
    
    def __init__(self, amber_model):
        """
        Initialize comparator with AMBER model
        
        Parameters:
        -----------
        amber_model : AMBERModel
            Trained AMBER model instance
        """
        self.amber_model = amber_model
        self.metrics_calc = MetricsCalculator()
        self.css_calc = ContextSensitivityScore()
    
    def find_most_similar_words(
        self, 
        target_vector: np.ndarray, 
        topn: int = 5,
        exclude_words: Optional[set] = None
    ) -> List[Tuple[str, float]]:
        """
        Find most similar words to a given vector
        
        Parameters:
        -----------
        target_vector : np.ndarray
            Target embedding vector
        topn : int, default=5
            Number of similar words to return
        exclude_words : set, optional
            Words to exclude from results
            
        Returns:
        --------
        List[Tuple[str, float]]
            List of (word, similarity_score) tuples
        """
        if exclude_words is None:
            exclude_words = set()
        
        similarities = []
        # Sample vocabulary for efficiency with large models
        vocab_sample = list(self.amber_model.w2v_model.key_to_index.keys())[:10000]
        
        for word in vocab_sample:
            if word not in exclude_words:
                word_vec = self.amber_model.w2v_model[word]
                if np.linalg.norm(target_vector) > 1e-8 and np.linalg.norm(word_vec) > 1e-8:
                    sim = 1 - cosine(target_vector, word_vec)
                    similarities.append((word, sim))
        
        similarities.sort(key=lambda x: x[1], reverse=True)
        return similarities[:topn]
    
    def evaluate_disambiguation(
        self, 
        test_cases: List[Dict],
        methods: List[str] = ["multi_head", "positional", "tfidf_only"]
    ) -> Dict[str, List]:
        """
        Evaluate different AMBER methods on word disambiguation tasks
        
        Parameters:
        -----------
        test_cases : List[Dict]
            Test cases with format:
            [{'word': str, 'contexts': [{'sentence': str, 'doc_idx': int, 'type': str}]}]
        methods : List[str]
            AMBER methods to evaluate
            
        Returns:
        --------
        Dict[str, List]
            Disambiguation results for each method
        """
        results = {method: [] for method in methods}
        results['word2vec_baseline'] = []
        
        for case in test_cases:
            target_word = case['word']
            contexts = case['contexts']
            
            print(f"\n=== Evaluating '{target_word}' ===")
            
            # Baseline Word2Vec (static)
            if target_word in self.amber_model.w2v_model.key_to_index:
                w2v_vec = self.amber_model.w2v_model[target_word]
                w2v_similar = self.find_most_similar_words(w2v_vec, exclude_words={target_word})
                results['word2vec_baseline'].append(w2v_similar)
                print(f"\nWord2Vec (Static): {[f'{w}({s:.3f})' for w, s in w2v_similar]}")
            
            # AMBER methods
            for method in methods:
                method_results = []
                
                for i, context_info in enumerate(contexts):
                    sentence = context_info['sentence']
                    doc_idx = context_info.get('doc_idx', i)
                    context_type = context_info.get('type', f'Context {i+1}')
                    
                    print(f"\n--- {context_type}: '{sentence}' ---")
                    
                    # Get contextual embedding
                    contextual_vec = self.amber_model.get_contextual_embedding(
                        target_word, sentence, method, doc_idx
                    )
                    
                    # Find similar words
                    similar_words = self.find_most_similar_words(
                        contextual_vec, exclude_words={target_word}
                    )
                    method_results.append(similar_words)
                    print(f"{method.replace('_', ' ').title()}: {[f'{w}({s:.3f})' for w, s in similar_words]}")
                
                results[method].append(method_results)
        
        return results
    
    def calculate_context_sensitivity(
        self, 
        word: str, 
        contexts: List[Dict],
        methods: List[str] = ["multi_head", "positional", "tfidf_only"]
    ) -> Dict[str, float]:
        """
        Calculate context sensitivity scores for different methods
        
        Parameters:
        -----------
        word : str
            Target word to analyze
        contexts : List[Dict]
            List of context dictionaries
        methods : List[str]
            Methods to evaluate
            
        Returns:
        --------
        Dict[str, float]
            Context sensitivity scores for each method
        """
        sensitivity_scores = {}
        
        # Baseline Word2Vec (should be 0)
        if word in self.amber_model.w2v_model.key_to_index:
            w2v_vec = self.amber_model.w2v_model[word]
            baseline_vecs = [w2v_vec] * len(contexts)
            sensitivity_scores['word2vec_baseline'] = self.css_calc.calculate(baseline_vecs)
        
        # AMBER methods
        for method in methods:
            method_vectors = []
            
            for i, context_info in enumerate(contexts):
                sentence = context_info['sentence']
                doc_idx = context_info.get('doc_idx', i)
                
                contextual_vec = self.amber_model.get_contextual_embedding(
                    word, sentence, method, doc_idx
                )
                method_vectors.append(contextual_vec)
            
            sensitivity_scores[method] = self.css_calc.calculate(method_vectors)
        
        return sensitivity_scores
    
    def comprehensive_comparison(
        self,
        test_cases: List[Dict],
        methods: List[str] = ["multi_head", "positional", "tfidf_only"],
        include_baseline: bool = True
    ) -> pd.DataFrame:
        """
        Comprehensive comparison across all test cases and methods
        
        Parameters:
        -----------
        test_cases : List[Dict]
            Test cases for evaluation
        methods : List[str]
            AMBER methods to evaluate
        include_baseline : bool, default=True
            Whether to include Word2Vec baseline
            
        Returns:
        --------
        pd.DataFrame
            Comprehensive comparison results
        """
        results_data = []
        
        all_methods = methods.copy()
        if include_baseline:
            all_methods.append('word2vec_baseline')
        
        for case in test_cases:
            word = case['word']
            contexts = case['contexts']
            
            # Calculate CSS for all methods
            css_scores = self.calculate_context_sensitivity(word, contexts, methods)
            
            # Calculate disambiguation accuracy (clustering purity)
            for method in all_methods:
                if method == 'word2vec_baseline':
                    if word in self.amber_model.w2v_model.key_to_index:
                        # Static embeddings
                        w2v_vec = self.amber_model.w2v_model[word]
                        embeddings = [w2v_vec] * len(contexts)
                    else:
                        embeddings = [np.zeros(self.amber_model.vector_size)] * len(contexts)
                else:
                    # AMBER embeddings
                    embeddings = []
                    for i, context_info in enumerate(contexts):
                        sentence = context_info['sentence']
                        doc_idx = context_info.get('doc_idx', i)
                        
                        embedding = self.amber_model.get_contextual_embedding(
                            word, sentence, method, doc_idx
                        )
                        embeddings.append(embedding)
                
                # Calculate metrics
                labels = list(range(len(contexts)))  # Each context is a different sense
                clustering_purity = self.metrics_calc.clustering_purity(embeddings, labels)
                css_score = css_scores.get(method, 0.0)
                
                results_data.append({
                    'word': word,
                    'method': method,
                    'css_score': css_score,
                    'clustering_purity': clustering_purity,
                    'num_contexts': len(contexts)
                })
        
        return pd.DataFrame(results_data)
    
    def visualize_comparison(
        self,
        comparison_df: pd.DataFrame,
        save_path: Optional[str] = None
    ) -> None:
        """
        Create visualization of comparison results
        
        Parameters:
        -----------
        comparison_df : pd.DataFrame
            Results from comprehensive_comparison
        save_path : str, optional
            Path to save the plot
        """
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        
        # CSS Score comparison
        axes[0, 0].set_title('Context Sensitivity Score (CSS) by Method')
        sns.boxplot(data=comparison_df, x='method', y='css_score', ax=axes[0, 0])
        axes[0, 0].set_xticklabels(axes[0, 0].get_xticklabels(), rotation=45)
        axes[0, 0].set_ylabel('CSS Score')
        
        # Clustering Purity comparison
        axes[0, 1].set_title('Clustering Purity by Method')
        sns.boxplot(data=comparison_df, x='method', y='clustering_purity', ax=axes[0, 1])
        axes[0, 1].set_xticklabels(axes[0, 1].get_xticklabels(), rotation=45)
        axes[0, 1].set_ylabel('Clustering Purity')
        
        # CSS by word
        axes[1, 0].set_title('CSS Score by Word')
        sns.boxplot(data=comparison_df, x='word', y='css_score', ax=axes[1, 0])
        axes[1, 0].set_xticklabels(axes[1, 0].get_xticklabels(), rotation=45)
        axes[1, 0].set_ylabel('CSS Score')
        
        # Method performance heatmap
        pivot_df = comparison_df.pivot_table(
            index='word', 
            columns='method', 
            values='css_score', 
            aggfunc='mean'
        )
        sns.heatmap(pivot_df, annot=True, fmt='.3f', cmap='YlOrRd', ax=axes[1, 1])
        axes[1, 1].set_title('CSS Score Heatmap')
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        
        plt.show()
    
    def generate_report(
        self,
        test_cases: List[Dict],
        methods: List[str] = ["multi_head", "positional", "tfidf_only"]
    ) -> str:
        """
        Generate comprehensive evaluation report
        
        Parameters:
        -----------
        test_cases : List[Dict]
            Test cases for evaluation
        methods : List[str]
            Methods to evaluate
            
        Returns:
        --------
        str
            Formatted evaluation report
        """
        # Run comprehensive comparison
        comparison_df = self.comprehensive_comparison(test_cases, methods)
        
        # Calculate summary statistics
        method_stats = comparison_df.groupby('method').agg({
            'css_score': ['mean', 'std', 'min', 'max'],
            'clustering_purity': ['mean', 'std', 'min', 'max']
        }).round(4)
        
        # Generate report
        report = []
        report.append("="*80)
        report.append("AMBER MODEL EVALUATION REPORT")
        report.append("="*80)
        report.append(f"\nCorpus Size: {len(self.amber_model.processed_corpus)} documents")
        
        if self.amber_model.tfidf_matrix is not None:
            report.append(f"TF-IDF Vocabulary: {len(self.amber_model.feature_names)} words")
        
        report.append(f"Word2Vec Model: {type(self.amber_model.w2v_model).__name__}")
        report.append(f"Vector Dimension: {self.amber_model.vector_size}")
        report.append(f"\nTest Cases: {len(test_cases)} polysemous words")
        report.append(f"Total Contexts: {sum(len(case['contexts']) for case in test_cases)}")
        
        report.append("\n" + "="*60)
        report.append("METHOD PERFORMANCE SUMMARY")
        report.append("="*60)
        
        for method in comparison_df['method'].unique():
            method_data = comparison_df[comparison_df['method'] == method]
            css_mean = method_data['css_score'].mean()
            css_std = method_data['css_score'].std()
            purity_mean = method_data['clustering_purity'].mean()
            purity_std = method_data['clustering_purity'].std()
            
            report.append(f"\n{method.replace('_', ' ').title()}:")
            report.append(f"  Context Sensitivity Score: {css_mean:.4f} ± {css_std:.4f}")
            report.append(f"  Clustering Purity: {purity_mean:.4f} ± {purity_std:.4f}")
        
        report.append("\n" + "="*60)
        report.append("WORD-SPECIFIC ANALYSIS")
        report.append("="*60)
        
        for word in comparison_df['word'].unique():
            word_data = comparison_df[comparison_df['word'] == word]
            report.append(f"\n'{word}':")
            
            for method in word_data['method'].unique():
                method_data = word_data[word_data['method'] == method]
                css = method_data['css_score'].iloc[0]
                purity = method_data['clustering_purity'].iloc[0]
                report.append(f"  {method}: CSS={css:.4f}, Purity={purity:.4f}")
        
        report.append("\n" + "="*60)
        report.append("CONCLUSIONS")
        report.append("="*60)
        
        # Find best performing method
        method_means = comparison_df.groupby('method')['css_score'].mean()
        best_method = method_means.idxmax()
        best_score = method_means.max()
        
        report.append(f"\n✓ Best performing method: {best_method.replace('_', ' ').title()}")
        report.append(f"✓ Achieved CSS: {best_score:.4f}")
        report.append(f"✓ Successfully demonstrates context-aware word disambiguation")
        
        return "\n".join(report)