"""Local TRACe-proxy metrics, as used by Notebook 5.

- **Context Relevance**: `src.reranking.reranker.relevance_score` (cross-encoder).
- **Adherence**: NLI entailment probability of context -> answer.
- **Completeness**: ROUGE-L F1 of answer vs. the gold reference response.
- **Utilization**: fraction of context sentences entailed by the answer.

`NliScorer` wraps a single `cross-encoder/nli-deberta-v3-base` model for both
adherence and utilization.
"""

import numpy as np
from rouge_score import rouge_scorer
from sentence_transformers import CrossEncoder

from src.reranking.reranker import split_sentences

NLI_MODEL_NAME = "cross-encoder/nli-deberta-v3-base"
NLI_LABELS = ["contradiction", "entailment", "neutral"]


def _softmax(x: np.ndarray) -> np.ndarray:
    e = np.exp(x - np.max(x))
    return e / e.sum()


class NliScorer:
    def __init__(self, model_name: str = NLI_MODEL_NAME):
        self.model = CrossEncoder(model_name)

    def entailment_prob(self, premise: str, hypothesis: str) -> float:
        logits = self.model.predict([(premise, hypothesis)])[0]
        probs = _softmax(np.array(logits))
        return float(probs[NLI_LABELS.index("entailment")])

    def adherence(self, context: str, answer: str) -> float:
        """P(context entails answer) -- is the answer grounded in the context?"""
        return self.entailment_prob(context, answer)

    def adherence_traced(self, context: str, answer: str) -> tuple[float, dict]:
        """Same as adherence() but also returns intermediate trace data."""
        prob = self.entailment_prob(context, answer)
        return prob, {"entailment_prob": prob}

    def utilization(self, context: str, answer: str, threshold: float = 0.5) -> float:
        """Fraction of context sentences entailed by the answer."""
        sentences = split_sentences(context)
        if not sentences:
            return 0.0
        pairs = [(answer, sentence) for sentence in sentences]
        logits = self.model.predict(pairs)
        entailed = sum(
            1 for logit in logits
            if _softmax(np.array(logit))[NLI_LABELS.index("entailment")] > threshold
        )
        return entailed / len(sentences)

    def utilization_traced(self, context: str, answer: str, threshold: float = 0.5) -> tuple[float, dict]:
        """Same as utilization() but also returns per-sentence entailment probabilities."""
        sentences = split_sentences(context)
        if not sentences:
            return 0.0, {"per_sentence": [], "threshold": threshold, "entailed_count": 0, "total": 0}
        pairs = [(answer, sentence) for sentence in sentences]
        logits = self.model.predict(pairs)
        per_sentence = []
        entailed_count = 0
        for sentence, logit in zip(sentences, logits):
            prob = float(_softmax(np.array(logit))[NLI_LABELS.index("entailment")])
            above = prob > threshold
            if above:
                entailed_count += 1
            per_sentence.append({"sentence": sentence, "entailment_prob": prob, "entailed": above})
        score = entailed_count / len(sentences)
        return score, {
            "per_sentence": per_sentence,
            "threshold": threshold,
            "entailed_count": entailed_count,
            "total": len(sentences),
        }


_rouge_scorer = rouge_scorer.RougeScorer(["rougeL"], use_stemmer=True)


def completeness(answer: str, reference: str) -> float:
    """ROUGE-L F1 of the generated answer vs. the gold reference response."""
    if not reference:
        return float("nan")
    return _rouge_scorer.score(reference, answer)["rougeL"].fmeasure