"""
Perplexity analysis module for measuring model uncertainty.
Mathematical formulation: PPL = exp(-1/N * Σ log p(w_i | w_{<i}))
"""

from typing import List, Optional
import math
import numpy as np
import torch


class PerplexityAnalyzer:
    """Computes perplexity of generated texts under the generating model."""
    
    def __init__(self, tokenizer, model):
        """
        Initialize perplexity analyzer.
        
        Args:
            tokenizer: Model tokenizer
            model: Language model for perplexity calculation
        """
        self.tokenizer = tokenizer
        self.model = model
    
    def compute_perplexity(self, texts: List[str]) -> Optional[float]:
        """
        Compute perplexity of generated texts.
        
        PPL = exp(-1/N * Σ log p(w_i | w_{<i}))
        
        Args:
            texts: List of generated texts
            
        Returns:
            Perplexity value (lower is usually better)
        """
        clean = self._filter_valid_texts(texts)
        if not clean:
            return None
        
        device = next(self.model.parameters()).device
        losses = []
        
        for text in clean:
            try:
                enc = self.tokenizer(
                    text, 
                    return_tensors="pt", 
                    truncation=True, 
                    max_length=512
                )
                input_ids = enc["input_ids"].to(device)
                
                if input_ids.shape[1] < 2:
                    continue
                
                with torch.no_grad():
                    outputs = self.model(input_ids=input_ids, labels=input_ids)
                
                if outputs.loss is not None and torch.isfinite(outputs.loss):
                    losses.append(float(outputs.loss.detach().cpu()))
            except Exception:
                continue
        
        if not losses:
            return None
        
        mean_loss = float(np.mean(losses))
        
        # Prevent overflow
        if mean_loss > 20:
            return float("inf")
        
        return float(math.exp(mean_loss))
    
    def compute_cross_entropy(self, texts: List[str]) -> Optional[float]:
        """
        Compute mean cross-entropy loss.
        
        Args:
            texts: List of generated texts
            
        Returns:
            Mean cross-entropy value
        """
        clean = self._filter_valid_texts(texts)
        if not clean:
            return None
        
        device = next(self.model.parameters()).device
        losses = []
        
        for text in clean:
            try:
                enc = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
                input_ids = enc["input_ids"].to(device)
                
                if input_ids.shape[1] < 2:
                    continue
                
                with torch.no_grad():
                    outputs = self.model(input_ids=input_ids, labels=input_ids)
                
                if outputs.loss is not None:
                    losses.append(float(outputs.loss.detach().cpu()))
            except Exception:
                continue
        
        return float(np.mean(losses)) if losses else None
    
    @staticmethod
    def _filter_valid_texts(texts: List[str]) -> List[str]:
        """Filter out empty or error texts."""
        return [t.strip() for t in texts if t and not t.startswith("ERROR")]
    
    def compute_all(self, texts: List[str]) -> dict:
        """Compute both perplexity and cross-entropy."""
        return {
            "Perplexity": self.compute_perplexity(texts),
            "Cross-Entropy": self.compute_cross_entropy(texts),
        }