"""
Quality analysis module.

Primary framework metric (proposal Equations xx-xxi):
(xx)  Q(y^{(m)}) = exp( (1/T_m) Σ_t log P_ref(x_t | x_<t) )
      Geometric-mean token likelihood under a FROZEN reference language
      model P_ref. Equal to the inverse perplexity of the sequence under
      the reference model. Because P_ref stays frozen, changes in Q̄
      reflect output changes rather than evaluator drift.
(xxi) Q̄ = 1/M Σ_m Q(y^{(m)})

Auxiliary evaluation metrics (proposal evaluation table): BERTScore, BLEU,
and a reference-free lexical fallback heuristic.
"""

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

try:
    from bert_score import score as bert_score
    BERTSCORE_AVAILABLE = True
except ImportError:
    BERTSCORE_AVAILABLE = False

try:
    import sacrebleu
    SACREBLEU_AVAILABLE = True
except ImportError:
    SACREBLEU_AVAILABLE = False


class FrozenReferenceQuality:
    """
    Semantic quality score under a frozen reference language model.

    Implements Equations (xx)-(xxi) of the proposal. The reference model
    P_ref is loaded once, kept in eval mode, and never updated, so the
    quality scale is stable across baseline and optimised generations.
    """

    def __init__(self, model_name: str = "distilgpt2", device: Optional[str] = None):
        """
        Args:
            model_name: Hugging Face id of the frozen reference LM. A small
                model (e.g. distilgpt2) keeps CPU evaluation tractable.
            device: Optional torch device override.
        """
        self.model_name = model_name
        self._device = device
        self._model = None
        self._tokenizer = None

    def _ensure_loaded(self) -> bool:
        if self._model is not None:
            return True
        try:
            import torch
            from transformers import AutoModelForCausalLM, AutoTokenizer

            self._torch = torch
            device = self._device or ("cuda" if torch.cuda.is_available() else "cpu")
            self._tokenizer = AutoTokenizer.from_pretrained(self.model_name)
            self._model = AutoModelForCausalLM.from_pretrained(self.model_name)
            self._model.to(device)
            self._model.eval()
            for p in self._model.parameters():  # frozen: no gradient updates
                p.requires_grad_(False)
            self._device = device
            return True
        except Exception:
            self._model = None
            return False

    def get_reference_model(self):
        """Expose the frozen P_ref (model, tokenizer), e.g. as a proxy
        scorer for logit-based metrics when the generating model is served
        behind an API and its logits are unavailable."""
        if not self._ensure_loaded():
            return None, None
        return self._model, self._tokenizer

    def compute_sequence_quality(self, text: str, max_length: int = 512) -> Optional[float]:
        """
        Compute Q(y^{(m)}) from Equation (xx).

        Q = exp( (1/T) Σ log P_ref(x_t | x_<t) ) = exp(-mean NLL) ∈ (0, 1]
        """
        if not text or text.startswith("ERROR"):
            return None
        if not self._ensure_loaded():
            return None

        torch = self._torch
        try:
            enc = self._tokenizer(
                text, return_tensors="pt", truncation=True, max_length=max_length
            )
            input_ids = enc["input_ids"].to(self._device)
            if input_ids.shape[1] < 2:
                return None

            with torch.no_grad():
                out = self._model(input_ids=input_ids, labels=input_ids)

            if out.loss is None or not torch.isfinite(out.loss):
                return None

            mean_nll = float(out.loss.detach().cpu())
            # Q = exp(mean log-prob) = exp(-mean NLL)
            return float(math.exp(-mean_nll))
        except Exception:
            return None

    def compute_aggregate_quality(self, texts: List[str]) -> Optional[float]:
        """
        Compute Q̄ from Equation (xxi): mean of per-sequence quality scores.
        """
        scores = []
        for text in texts:
            q = self.compute_sequence_quality(text)
            if q is not None:
                scores.append(q)
        if not scores:
            return None
        return float(np.mean(scores))

    # ------------------------------------------------------------------
    # Calibration for threshold comparison
    #
    # Raw Q from Eq (xx) is a geometric-mean token probability, i.e. the
    # inverse perplexity of the sequence: fluent English under a small
    # reference LM typically lands around Q ≈ 0.02-0.05. The dynamic
    # quality thresholds Q_th(x) (Eq xxviii) are defined on a [0,1] score
    # scale. To compare the two we apply the monotone, order-preserving
    # calibration
    #
    #     Q_cal = Q_raw^(1/τ_q) = exp( (1/τ_q) · mean log P_ref )
    #
    # which is a temperature-scaled geometric mean (τ_q = 10 by default:
    # mean NLL 2.0 → 0.82, 3.5 → 0.70, 5.0 → 0.61). The raw Eq (xx) value
    # is always reported alongside.
    # ------------------------------------------------------------------

    CALIBRATION_TAU = 10.0

    def compute_calibrated_quality(self, text: str) -> Optional[float]:
        """Calibrated per-sequence quality Q_cal = Q_raw^(1/τ_q)."""
        q = self.compute_sequence_quality(text)
        if q is None or q <= 0:
            return None
        return float(q ** (1.0 / self.CALIBRATION_TAU))

    def compute_aggregate_calibrated_quality(self, texts: List[str]) -> Optional[float]:
        """Aggregate calibrated quality (Eq xxi applied to Q_cal)."""
        scores = []
        for text in texts:
            q = self.compute_calibrated_quality(text)
            if q is not None:
                scores.append(q)
        if not scores:
            return None
        return float(np.mean(scores))


class QualityAnalyzer:
    """Analyzes semantic and lexical quality against reference answers."""
    
    def __init__(self, bertscore_model: str = "distilbert-base-uncased"):
        """
        Initialize quality analyzer.
        
        Args:
            bertscore_model: Model name for BERTScore
        """
        self.bertscore_model = bertscore_model
    
    def compute_bertscore(self, texts: List[str], reference: str) -> Optional[float]:
        """
        Compute BERTScore F1 between generated texts and reference.
        
        Higher score indicates better semantic similarity.
        
        Args:
            texts: List of generated texts
            reference: Reference answer string
            
        Returns:
            Mean BERTScore F1 (0-1), higher is better
        """
        if not BERTSCORE_AVAILABLE:
            return None
        
        clean = self._filter_valid_texts(texts)
        ref = (reference or "").strip()
        
        if not clean or not ref:
            return None
        
        try:
            _, _, f1 = bert_score(
                clean,
                [ref] * len(clean),
                lang="en",
                model_type=self.bertscore_model,
                verbose=False,
                rescale_with_baseline=False,
            )
            return float(f1.mean())
        except Exception:
            return None
    
    def compute_bleu(self, texts: List[str], reference: str) -> Optional[float]:
        """
        Compute BLEU score between generated texts and reference.
        
        Higher score indicates better n-gram overlap.
        
        Args:
            texts: List of generated texts
            reference: Reference answer string
            
        Returns:
            BLEU score (0-100), higher is better
        """
        if not SACREBLEU_AVAILABLE:
            return None
        
        clean = self._filter_valid_texts(texts)
        ref = (reference or "").strip()
        
        if not clean or not ref:
            return None
        
        try:
            bleu = sacrebleu.corpus_bleu(clean, [[ref] * len(clean)])
            return float(bleu.score)
        except Exception:
            return None
    
    def compute_fallback_quality(self, texts: List[str]) -> float:
        """
        Compute reference-free quality heuristic using lexical features.
        
        Used when no reference answer is supplied.
        
        Args:
            texts: List of generated texts
            
        Returns:
            Quality estimate (0-1), higher is better
        """
        clean = self._filter_valid_texts(texts)
        if not clean:
            return 0.0
        
        combined = " ".join(clean)
        
        # Average word length (proxy for vocabulary sophistication)
        words = combined.split()
        avg_word_len = min(1.0, sum(len(w) for w in words) / max(1, len(words)) / 10)
        
        # Sentence length variability
        sentences = self._split_sentences(combined)
        if len(sentences) > 1:
            lengths = [len(s.split()) for s in sentences]
            length_variability = min(1.0, np.std(lengths) / 10)
        else:
            length_variability = 0.5
        
        # Presence of punctuation variety
        punctuation_chars = set(".!?;:,'\"-")
        punct_count = sum(1 for c in combined if c in punctuation_chars)
        punct_density = min(1.0, punct_count / max(1, len(combined)) * 20)
        
        quality = (avg_word_len * 0.3 + length_variability * 0.35 + punct_density * 0.35)
        return float(min(1.0, max(0.0, quality)))
    
    @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")]
    
    @staticmethod
    def _split_sentences(text: str) -> List[str]:
        """Split text into sentences."""
        import re
        sentences = re.split(r'[.!?]+', text)
        return [s.strip() for s in sentences if len(s.strip()) > 3]
    
    def compute_quality_metrics(
        self, 
        texts: List[str], 
        reference: Optional[str] = None
    ) -> dict:
        """Compute all quality metrics."""
        metrics = {
            "Fallback Quality": self.compute_fallback_quality(texts),
        }
        
        if reference:
            metrics["BERTScore F1"] = self.compute_bertscore(texts, reference)
            metrics["BLEU"] = self.compute_bleu(texts, reference)
        
        return metrics