comparison_metrics.py

"""
Comparison metrics for measuring information loss (Objective 4 / RQ3).

Compares topic models trained on full text vs summarized text.
"""
import numpy as np
import pandas as pd
from typing import List, Dict, Tuple, Set
from bertopic import BERTopic
from sklearn.metrics.pairwise import cosine_similarity
from scipy.stats import spearmanr, kendalltau
from scipy.spatial.distance import jensenshannon
import warnings
warnings.filterwarnings('ignore')


class TopicComparisonMetrics:
    """
    Metrics for comparing two topic models (full text vs summary).
    """

    def __init__(self, full_model: BERTopic, summary_model: BERTopic,
                 full_docs: List[str], summary_docs: List[str]):
        """
        Initialize comparison metrics.

        Parameters:
        -----------
        full_model : BERTopic
            Topic model trained on full text
        summary_model : BERTopic
            Topic model trained on summaries
        full_docs : List[str]
            Full text documents
        summary_docs : List[str]
            Summary documents
        """
        self.full_model = full_model
        self.summary_model = summary_model
        self.full_docs = full_docs
        self.summary_docs = summary_docs

    def calculate_topic_count_diff(self) -> Dict[str, int]:
        """
        Compare number of topics detected.

        Returns:
        --------
        Dict[str, int]
            Topic counts and difference
        """
        full_topics = [t for t in set(self.full_model.topics_) if t != -1]
        summary_topics = [t for t in set(self.summary_model.topics_) if t != -1]

        # Cap ratio at 1.0 - having more topics doesn't mean better preservation
        # It could indicate overfitting or fragmentation
        raw_ratio = len(summary_topics) / len(full_topics) if full_topics else 0
        capped_ratio = min(raw_ratio, 1.0)

        return {
            'full_text_topics': len(full_topics),
            'summary_topics': len(summary_topics),
            'difference': len(full_topics) - len(summary_topics),
            'ratio': capped_ratio,
            'raw_ratio': raw_ratio  # Keep raw ratio for analysis
        }

    def calculate_topic_overlap(self, similarity_threshold: float = 0.7) -> Dict[str, float]:
        """
        Calculate topic overlap based on keyword similarity.

        Parameters:
        -----------
        similarity_threshold : float
            Threshold for considering topics as "same"

        Returns:
        --------
        Dict[str, float]
            Overlap metrics
        """
        full_topics = self.full_model.get_topics()
        summary_topics = self.summary_model.get_topics()

        # Remove outlier topics
        full_topics = {k: v for k, v in full_topics.items() if k != -1}
        summary_topics = {k: v for k, v in summary_topics.items() if k != -1}

        if not full_topics or not summary_topics:
            return {'overlap_count': 0, 'overlap_ratio': 0.0, 'similarity_scores': []}

        # Get top words for each topic
        full_words_dict = {
            k: set([word for word, _ in v[:20]])
            for k, v in full_topics.items()
        }

        summary_words_dict = {
            k: set([word for word, _ in v[:20]])
            for k, v in summary_topics.items()
        }

        # Calculate Jaccard similarity between all topic pairs
        matches = 0
        all_similarities = []

        for full_id, full_words in full_words_dict.items():
            max_similarity = 0

            for summary_id, summary_words in summary_words_dict.items():
                # Jaccard similarity
                intersection = len(full_words & summary_words)
                union = len(full_words | summary_words)
                similarity = intersection / union if union > 0 else 0

                all_similarities.append(similarity)
                max_similarity = max(max_similarity, similarity)

            if max_similarity >= similarity_threshold:
                matches += 1

        overlap_ratio = matches / len(full_topics) if full_topics else 0

        return {
            'overlap_count': matches,
            'overlap_ratio': overlap_ratio,
            'avg_similarity': np.mean(all_similarities) if all_similarities else 0,
            'max_similarity': np.max(all_similarities) if all_similarities else 0
        }

    def calculate_vocabulary_preservation(self) -> Dict[str, float]:
        """
        Calculate how much of the important vocabulary is preserved.

        Returns:
        --------
        Dict[str, float]
            Vocabulary preservation metrics
        """
        # Get all important words from full text topics
        full_topics = self.full_model.get_topics()
        summary_topics = self.summary_model.get_topics()

        full_words = set()
        for topic_id, words in full_topics.items():
            if topic_id != -1:
                full_words.update([word for word, _ in words[:20]])

        summary_words = set()
        for topic_id, words in summary_topics.items():
            if topic_id != -1:
                summary_words.update([word for word, _ in words[:20]])

        if not full_words:
            return {'preserved_ratio': 0.0, 'lost_words': 0}

        preserved = full_words & summary_words
        lost = full_words - summary_words

        return {
            'preserved_ratio': len(preserved) / len(full_words),
            'preserved_count': len(preserved),
            'lost_count': len(lost),
            'new_words': len(summary_words - full_words)
        }

    def calculate_document_assignment_consistency(self) -> Dict[str, float]:
        """
        Calculate consistency of document-to-topic assignments.

        Checks if documents are assigned to "similar" topics in both models.

        Returns:
        --------
        Dict[str, float]
            Assignment consistency metrics
        """
        full_assignments = self.full_model.topics_
        summary_assignments = self.summary_model.topics_

        if len(full_assignments) != len(summary_assignments):
            print("Warning: Document counts differ between models")
            return {'consistency': 0.0}

        # Build topic similarity matrix
        full_topics = self.full_model.get_topics()
        summary_topics = self.summary_model.get_topics()

        # For each document, check if assigned to "similar" topics
        consistent_assignments = 0

        for i in range(len(full_assignments)):
            full_topic = full_assignments[i]
            summary_topic = summary_assignments[i]

            # Skip if either is outlier
            if full_topic == -1 or summary_topic == -1:
                continue

            # Get topic words
            if full_topic in full_topics and summary_topic in summary_topics:
                full_words = set([word for word, _ in full_topics[full_topic][:20]])
                summary_words = set([word for word, _ in summary_topics[summary_topic][:20]])

                # Calculate similarity
                intersection = len(full_words & summary_words)
                union = len(full_words | summary_words)
                similarity = intersection / union if union > 0 else 0

                if similarity >= 0.3:  # Lower threshold for individual docs
                    consistent_assignments += 1

        total_docs = len([t for t in full_assignments if t != -1])
        consistency = consistent_assignments / total_docs if total_docs > 0 else 0

        return {
            'consistency': consistency,
            'consistent_docs': consistent_assignments,
            'total_docs': total_docs
        }

    def calculate_emerging_topics_preservation(self, full_emerging: pd.DataFrame,
                                              summary_emerging: pd.DataFrame) -> Dict[str, float]:
        """
        Calculate how well emerging topics are preserved.

        Parameters:
        -----------
        full_emerging : pd.DataFrame
            Emerging topics from full text
        summary_emerging : pd.DataFrame
            Emerging topics from summaries

        Returns:
        --------
        Dict[str, float]
            Preservation metrics
        """
        if full_emerging.empty or summary_emerging.empty:
            return {
                'emerging_overlap_count': 0,
                'emerging_overlap_ratio': 0.0,
                'emergence_score_correlation': 0.0
            }

        full_topics_set = set(full_emerging['Topic'].values)
        summary_topics_set = set(summary_emerging['Topic'].values)

        # Direct ID overlap (may not be meaningful if topics renumbered)
        direct_overlap = len(full_topics_set & summary_topics_set)

        # Compare emergence scores correlation
        if len(full_emerging) > 0 and len(summary_emerging) > 0:
            # Match topics by similarity (simplified - assumes same IDs for now)
            common_topics = list(full_topics_set & summary_topics_set)

            if common_topics:
                full_scores = full_emerging[full_emerging['Topic'].isin(common_topics)]['Emergence_Score'].values
                summary_scores = summary_emerging[summary_emerging['Topic'].isin(common_topics)]['Emergence_Score'].values

                if len(full_scores) == len(summary_scores) and len(full_scores) > 1:
                    correlation, _ = spearmanr(full_scores, summary_scores)
                else:
                    correlation = 0
            else:
                correlation = 0
        else:
            correlation = 0

        # Cap emerging overlap ratio at 1.0 (can't preserve more than 100%)
        raw_overlap_ratio = direct_overlap / len(full_topics_set) if full_topics_set else 0
        capped_overlap_ratio = min(raw_overlap_ratio, 1.0)

        return {
            'emerging_overlap_count': direct_overlap,
            'emerging_overlap_ratio': capped_overlap_ratio,
            'emergence_score_correlation': correlation,
            'raw_overlap_ratio': raw_overlap_ratio  # Keep for analysis
        }

    def calculate_burst_preservation(self, full_bursts: pd.DataFrame,
                                    summary_bursts: pd.DataFrame) -> Dict[str, float]:
        """
        Calculate how well burst patterns are preserved.

        Parameters:
        -----------
        full_bursts : pd.DataFrame
            Burst summary from full text
        summary_bursts : pd.DataFrame
            Burst summary from summaries

        Returns:
        --------
        Dict[str, float]
            Burst preservation metrics
        """
        if full_bursts.empty or summary_bursts.empty or len(full_bursts) == 0 or len(summary_bursts) == 0:
            return {
                'burst_topic_overlap': 0.0,
                'burst_count_correlation': 0.0
            }

        full_topics = set(full_bursts['Topic'].values)
        summary_topics = set(summary_bursts['Topic'].values)

        overlap = len(full_topics & summary_topics)

        # Compare consensus burst counts for common topics
        common_topics = list(full_topics & summary_topics)

        if common_topics:
            full_consensus = full_bursts[full_bursts['Topic'].isin(common_topics)]['Consensus_Bursts'].values
            summary_consensus = summary_bursts[summary_bursts['Topic'].isin(common_topics)]['Consensus_Bursts'].values

            if len(full_consensus) == len(summary_consensus) and len(full_consensus) > 0:
                # Correlation of burst counts
                if len(full_consensus) > 1:
                    correlation, _ = spearmanr(full_consensus, summary_consensus)
                else:
                    correlation = 1.0 if full_consensus[0] == summary_consensus[0] else 0.0
            else:
                correlation = 0
        else:
            correlation = 0

        # Cap burst overlap at 1.0 (can't preserve more than 100%)
        raw_burst_overlap = overlap / len(full_topics) if full_topics else 0
        capped_burst_overlap = min(raw_burst_overlap, 1.0)

        return {
            'burst_topic_overlap': capped_burst_overlap,
            'burst_count_correlation': correlation,
            'raw_burst_overlap': raw_burst_overlap  # Keep for analysis
        }

    def calculate_information_loss_score(self, full_coherence: float,
                                         summary_coherence: float,
                                         similarity_threshold: float = 0.7) -> Dict[str, float]:
        """
        Calculate comprehensive information loss score.

        Parameters:
        -----------
        full_coherence : float
            Coherence score for full text model
        summary_coherence : float
            Coherence score for summary model
        similarity_threshold : float
            Threshold for topic matching

        Returns:
        --------
        Dict[str, float]
            Comprehensive loss metrics
        """
        # Topic count preservation
        topic_diff = self.calculate_topic_count_diff()
        topic_count_preservation = topic_diff['ratio']

        # Topic overlap
        overlap = self.calculate_topic_overlap(similarity_threshold)
        topic_overlap_score = overlap['overlap_ratio']

        # Vocabulary preservation
        vocab = self.calculate_vocabulary_preservation()
        vocab_preservation = vocab['preserved_ratio']

        # Coherence preservation (cap at 1.0 - higher coherence doesn't mean better preservation)
        raw_coherence_ratio = summary_coherence / full_coherence if full_coherence > 0 else 0
        coherence_preservation = min(raw_coherence_ratio, 1.0)

        # Overall information loss (0 = total loss, 1 = no loss)
        preservation_score = (
            0.3 * topic_count_preservation +
            0.3 * topic_overlap_score +
            0.2 * vocab_preservation +
            0.2 * coherence_preservation
        )

        information_loss = 1.0 - preservation_score

        return {
            'information_loss': information_loss,
            'preservation_score': preservation_score,
            'topic_count_preservation': topic_count_preservation,
            'topic_overlap_score': topic_overlap_score,
            'vocab_preservation': vocab_preservation,
            'coherence_preservation': coherence_preservation
        }

    def generate_comparison_report(self, full_coherence: float,
                                   summary_coherence: float,
                                   full_emerging: pd.DataFrame = None,
                                   summary_emerging: pd.DataFrame = None,
                                   full_bursts: pd.DataFrame = None,
                                   summary_bursts: pd.DataFrame = None) -> Dict:
        """
        Generate comprehensive comparison report.

        Returns:
        --------
        Dict
            Complete comparison metrics
        """
        report = {
            'topic_counts': self.calculate_topic_count_diff(),
            'topic_overlap': self.calculate_topic_overlap(),
            'vocabulary': self.calculate_vocabulary_preservation(),
            'document_consistency': self.calculate_document_assignment_consistency(),
            'coherence': {
                'full_text': full_coherence,
                'summary': summary_coherence,
                'difference': full_coherence - summary_coherence,
                'ratio': summary_coherence / full_coherence if full_coherence > 0 else 0
            },
            'information_loss': self.calculate_information_loss_score(
                full_coherence, summary_coherence
            )
        }

        if (full_emerging is not None and summary_emerging is not None and
            not full_emerging.empty and not summary_emerging.empty):
            report['emerging_topics'] = self.calculate_emerging_topics_preservation(
                full_emerging, summary_emerging
            )

        if (full_bursts is not None and summary_bursts is not None and
            not full_bursts.empty and not summary_bursts.empty):
            report['bursts'] = self.calculate_burst_preservation(
                full_bursts, summary_bursts
            )

        return report