"""Semantic similarity scorers using embeddings."""

import numpy as np
import litellm

from openevalkit.scorers.base import Scorer
from openevalkit.score import Score
from openevalkit.run import Run
from openevalkit.errors import ScoringError


class CosineSimilarity(Scorer):
    """
    Semantic similarity between output and reference using embeddings.

    Uses ``litellm.embedding()`` to get embeddings from any supported
    provider (OpenAI, Cohere, etc.) and computes cosine similarity.

    Returns a float in [0.0, 1.0] where 1.0 means semantically identical.

    Examples:
        >>> scorer = CosineSimilarity(model="text-embedding-3-small")
        >>> run = Run(id="1", input="q", output="The cat sat", reference="A cat was sitting")
        >>> score = scorer.score(run)  # Requires API key
        >>> 0.0 <= score.value <= 1.0
        True
    """

    name = "cosine_similarity"
    requires_reference = True
    cacheable = True

    def __init__(self, model: str = "text-embedding-3-small", **kwargs):
        """
        Args:
            model: Embedding model name (any model supported by litellm).
            **kwargs: Additional keyword arguments passed to ``litellm.embedding()``.
        """
        self.model = model
        self.extra_kwargs = kwargs

    def score(self, run: Run) -> Score:
        if run.reference is None:
            raise ScoringError(
                f"{self.name} requires reference field. "
                f"Run {run.id} has no reference."
            )

        output_str = str(run.output)
        reference_str = str(run.reference)

        try:
            response = litellm.embedding(
                model=self.model,
                input=[output_str, reference_str],
                **self.extra_kwargs,
            )
            output_emb = np.array(response.data[0]["embedding"])
            reference_emb = np.array(response.data[1]["embedding"])
        except Exception as e:
            raise ScoringError(
                f"Embedding call failed for run {run.id} with model {self.model}: {e}"
            ) from e

        # Cosine similarity
        dot = np.dot(output_emb, reference_emb)
        norm_a = np.linalg.norm(output_emb)
        norm_b = np.linalg.norm(reference_emb)

        if norm_a == 0 or norm_b == 0:
            similarity = 0.0
        else:
            similarity = float(dot / (norm_a * norm_b))

        return Score(
            value=similarity,
            reason=f"Cosine similarity: {similarity:.4f}",
            metadata={
                "model": self.model,
                "embedding_dimensions": len(output_emb),
            },
        )
