"""Semantic entropy scorer — measures answer diversity as hallucination risk.

Clusters answers by cosine similarity and computes entropy over cluster counts.
High entropy = high disagreement = high risk.
Purpose:    Clusters answers by cosine similarity and computes Shannon entropy over cluster sizes. High entropy indicates high disagreement and elevated failure risk.
Exports:    SemanticEntropyScorer, SemanticEntropyResult
Added:      v0.1.0
Status:     Experimental
"""
from __future__ import annotations
import math
from dataclasses import dataclass
from typing import Optional


@dataclass
class SemanticEntropyResult:
    risk: float
    n_clusters: int
    entropy: float


class SemanticEntropyScorer:
    """Zero-label semantic diversity scorer for multi-answer chains.

    Args:
        model_name: sentence-transformers model.
        threshold: cosine distance above which answers are in different clusters.
        experimental: scorer is considered experimental.
    """

    def __init__(
        self,
        model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
        threshold: float = 0.3,
        experimental: bool = True,
    ) -> None:
        self.model_name = model_name
        self.threshold = threshold
        self.experimental = experimental
        self._model = None

    def _get_model(self):
        if self._model is None:
            from sentence_transformers import SentenceTransformer
            self._model = SentenceTransformer(self.model_name)
        return self._model

    def score(self, answers: list[str]) -> float:
        """Return risk in [0, 1]. High = high semantic entropy (disagreement)."""
        if len(answers) <= 1:
            return 0.5
        return self.score_detailed(answers).risk

    def score_detailed(self, answers: list[str]) -> SemanticEntropyResult:
        """Return SemanticEntropyResult with risk, n_clusters, entropy."""
        if len(answers) <= 1:
            return SemanticEntropyResult(risk=0.5, n_clusters=1, entropy=0.0)

        model = self._get_model()
        embeddings = model.encode(answers, normalize_embeddings=True)

        # Greedy clustering by cosine similarity (dot product for normalized vecs)
        clusters: list[list[int]] = []
        for i, emb_i in enumerate(embeddings):
            placed = False
            for c_idx, cluster in enumerate(clusters):
                rep = embeddings[cluster[0]]
                sim = float(emb_i @ rep)
                if sim >= (1.0 - self.threshold):
                    cluster.append(i)
                    placed = True
                    break
            if not placed:
                clusters.append([i])

        n = len(answers)
        counts = [len(c) for c in clusters]
        entropy = -sum((cnt / n) * math.log(cnt / n) for cnt in counts if cnt > 0)
        max_entropy = math.log(n) if n > 1 else 1.0
        risk = max(0.0, min(1.0, entropy / max_entropy)) if max_entropy > 0 else 0.0

        return SemanticEntropyResult(risk=risk, n_clusters=len(clusters), entropy=entropy)
