evaluation_metrics.py

"""
Evaluation metrics for topic modeling quality assessment (Objective 3).

This module provides comprehensive evaluation metrics including:
- Coherence scores (c_v, c_npmi, c_uci)
- Topic diversity metrics
- Model quality metrics
"""
import numpy as np
import pandas as pd
from typing import List, Dict, Tuple
from bertopic import BERTopic
from gensim.corpora import Dictionary
from gensim.models.coherencemodel import CoherenceModel
from sklearn.metrics.pairwise import cosine_similarity
import warnings
warnings.filterwarnings('ignore')


class TopicModelEvaluator:
    """
    Comprehensive evaluation of topic models.
    """

    def __init__(self, topic_model: BERTopic, docs: List[str], embeddings: np.ndarray = None):
        """
        Initialize evaluator.

        Parameters:
        -----------
        topic_model : BERTopic
            Trained BERTopic model
        docs : List[str]
            List of documents
        embeddings : np.ndarray
            Document embeddings (optional)
        """
        self.topic_model = topic_model
        self.docs = docs
        self.embeddings = embeddings
        self.topics = topic_model.topics_

    def calculate_coherence(self, coherence_type: str = 'c_v') -> float:
        """
        Calculate topic coherence score.

        Parameters:
        -----------
        coherence_type : str
            Type of coherence ('c_v', 'c_npmi', 'c_uci', 'u_mass')

        Returns:
        --------
        float
            Coherence score
        """
        try:
            # Get topics (excluding outliers)
            topics = self.topic_model.get_topics()
            topic_words = []

            for topic_id in sorted(topics.keys()):
                if topic_id not in (-1, 0):  # Exclude outlier topic and catch-all topic 0
                    words = [word for word, _ in topics[topic_id][:10]]
                    topic_words.append(words)

            if not topic_words:
                return 0.0

            # Tokenize documents
            texts = [doc.lower().split() for doc in self.docs]

            # Create dictionary and corpus
            dictionary = Dictionary(texts)

            # Calculate coherence
            coherence_model = CoherenceModel(
                topics=topic_words,
                texts=texts,
                dictionary=dictionary,
                coherence=coherence_type
            )

            return coherence_model.get_coherence()

        except Exception as e:
            print(f"Error calculating {coherence_type} coherence: {e}")
            return 0.0

    def calculate_all_coherence_scores(self) -> Dict[str, float]:
        """
        Calculate all coherence metrics.

        Returns:
        --------
        Dict[str, float]
            Dictionary with coherence scores
        """
        scores = {}

        for coherence_type in ['c_v', 'c_npmi', 'c_uci', 'u_mass']:
            print(f"Calculating {coherence_type} coherence...")
            scores[coherence_type] = self.calculate_coherence(coherence_type)

        return scores

    def calculate_topic_diversity(self) -> float:
        """
        Calculate topic diversity (uniqueness of topics).

        Higher diversity = more unique topics

        Returns:
        --------
        float
            Topic diversity score (0-1)
        """
        topics = self.topic_model.get_topics()

        if len(topics) <= 1:
            return 0.0

        # Get top words for each topic
        topic_words = []
        for topic_id in sorted(topics.keys()):
            if topic_id not in (-1, 0):
                words = [word for word, _ in topics[topic_id][:10]]
                topic_words.append(set(words))

        if len(topic_words) < 2:
            return 0.0

        # Calculate unique words across all topics
        all_words = set()
        for words in topic_words:
            all_words.update(words)

        # Calculate total words used
        total_words = sum(len(words) for words in topic_words)

        # Diversity = unique words / total words
        diversity = len(all_words) / total_words if total_words > 0 else 0.0

        return diversity

    def calculate_within_topic_diversity(self) -> float:
        """
        Calculate average diversity within each topic.

        Measures how spread out the word probabilities are within topics.

        Returns:
        --------
        float
            Average within-topic diversity
        """
        topics = self.topic_model.get_topics()
        diversities = []

        for topic_id in sorted(topics.keys()):
            if topic_id not in (-1, 0):
                # Get word probabilities
                word_probs = [prob for _, prob in topics[topic_id][:20]]

                if word_probs:
                    # Calculate entropy as diversity measure
                    probs = np.array(word_probs)
                    probs = probs / probs.sum()  # Normalize
                    entropy = -np.sum(probs * np.log(probs + 1e-10))
                    diversities.append(entropy)

        return np.mean(diversities) if diversities else 0.0

    def calculate_topic_significance(self) -> Dict[int, float]:
        """
        Calculate significance score for each topic.

        Based on number of documents and avg probability.

        Returns:
        --------
        Dict[int, float]
            Topic ID to significance score mapping
        """
        topics = self.topics
        unique_topics = set(topics)

        significance = {}

        for topic_id in unique_topics:
            if topic_id not in (-1, 0):
                # Count documents in topic
                doc_count = sum(1 for t in topics if t == topic_id)

                # Calculate relative size
                rel_size = doc_count / len(topics)

                significance[topic_id] = rel_size

        return significance

    def calculate_topic_quality_score(self) -> float:
        """
        Calculate overall topic quality score (0-100).

        Combines multiple metrics into a single score.

        Returns:
        --------
        float
            Overall quality score
        """
        try:
            # Get coherence (c_v is most interpretable)
            coherence = self.calculate_coherence('c_v')

            # Get diversity
            diversity = self.calculate_topic_diversity()

            # Get number of topics (penalize too few or too many)
            n_topics = len([t for t in set(self.topics) if t not in (-1, 0)])
            topic_count_score = min(n_topics / 20, 1.0)  # Optimal around 20

            # Combine scores (weighted average)
            quality = (
                0.5 * max(0, min(coherence, 1.0)) +  # Coherence (50%)
                0.3 * diversity +  # Diversity (30%)
                0.2 * topic_count_score  # Topic count (20%)
            ) * 100

            return quality

        except Exception as e:
            print(f"Error calculating quality score: {e}")
            return 0.0

    def get_topic_sizes(self) -> pd.DataFrame:
        """
        Get size distribution of topics.

        Returns:
        --------
        pd.DataFrame
            Topic sizes sorted by frequency
        """
        topic_info = self.topic_model.get_topic_info()
        return topic_info[['Topic', 'Count']].sort_values('Count', ascending=False)

    def get_comprehensive_metrics(self) -> Dict[str, any]:
        """
        Calculate all evaluation metrics at once.

        Returns:
        --------
        Dict[str, any]
            Dictionary with all metrics
        """
        print("Calculating comprehensive evaluation metrics...")

        metrics = {
            'n_topics': len([t for t in set(self.topics) if t not in (-1, 0)]),
            'n_outliers': sum(1 for t in self.topics if t in (-1, 0)),
            'coherence_scores': self.calculate_all_coherence_scores(),
            'topic_diversity': self.calculate_topic_diversity(),
            'within_topic_diversity': self.calculate_within_topic_diversity(),
            'quality_score': self.calculate_topic_quality_score(),
            'topic_significance': self.calculate_topic_significance(),
        }

        return metrics

    def print_evaluation_report(self):
        """
        Print a formatted evaluation report.
        """
        metrics = self.get_comprehensive_metrics()

        print("\n" + "="*80)
        print("TOPIC MODEL EVALUATION REPORT")
        print("="*80)

        print(f"\n📊 Basic Statistics:")
        print(f"  - Number of topics: {metrics['n_topics']}")
        print(f"  - Number of outliers: {metrics['n_outliers']}")
        print(f"  - Outlier percentage: {metrics['n_outliers']/len(self.docs)*100:.2f}%")

        print(f"\n🎯 Coherence Scores:")
        for coh_type, score in metrics['coherence_scores'].items():
            print(f"  - {coh_type.upper()}: {score:.4f}")

        print(f"\n🌈 Diversity Metrics:")
        print(f"  - Topic diversity: {metrics['topic_diversity']:.4f}")
        print(f"  - Within-topic diversity: {metrics['within_topic_diversity']:.4f}")

        print(f"\n⭐ Overall Quality Score: {metrics['quality_score']:.2f}/100")

        print("\n" + "="*80)

    def save_metrics_to_file(self, filepath: str):
        """
        Save evaluation metrics to JSON file.

        Parameters:
        -----------
        filepath : str
            Path to save metrics
        """
        import json

        metrics = self.get_comprehensive_metrics()

        # Convert numpy types to Python types
        def convert_types(obj):
            if isinstance(obj, np.integer):
                return int(obj)
            elif isinstance(obj, np.floating):
                return float(obj)
            elif isinstance(obj, dict):
                return {k: convert_types(v) for k, v in obj.items()}
            elif isinstance(obj, list):
                return [convert_types(item) for item in obj]
            return obj

        metrics = convert_types(metrics)

        with open(filepath, 'w') as f:
            json.dump(metrics, f, indent=2)

        print(f"Metrics saved to {filepath}")