from __future__ import annotations

import logging
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    List,
    Literal,
    Optional,
    Sequence,
    Union,
)

from typing_extensions import TypedDict

if TYPE_CHECKING:
    import numpy as np  # type: ignore


logger = logging.getLogger(__name__)

Matrix = Union[List[List[float]], List[Any], Any]


def cosine_similarity(X: Matrix, Y: Matrix) -> np.ndarray:
    """Row-wise cosine similarity between two equal-width matrices."""
    import numpy as np

    if len(X) == 0 or len(Y) == 0:
        return np.array([])

    X = np.array(X)
    Y = np.array(Y)
    if X.shape[1] != Y.shape[1]:
        raise ValueError(
            f"Number of columns in X and Y must be the same. X has shape {X.shape} "
            f"and Y has shape {Y.shape}."
        )
    try:
        import simsimd as simd  # type: ignore

        X = np.array(X, dtype=np.float32)
        Y = np.array(Y, dtype=np.float32)
        Z = 1 - simd.cdist(X, Y, metric="cosine")
        if isinstance(Z, float):
            return np.array([Z])
        return np.array(Z)
    except ImportError:
        logger.debug(
            "Unable to import simsimd, defaulting to NumPy implementation. If you want "
            "to use simsimd please install with `pip install simsimd`."
        )
        X_norm = np.linalg.norm(X, axis=1)
        Y_norm = np.linalg.norm(Y, axis=1)
        # Ignore divide by zero errors run time warnings as those are handled below.
        with np.errstate(divide="ignore", invalid="ignore"):
            similarity = np.dot(X, Y.T) / np.outer(X_norm, Y_norm)
        similarity[np.isnan(similarity) | np.isinf(similarity)] = 0.0
        return similarity


def _get_openai_encoder() -> Callable[[Sequence[str]], Sequence[Sequence[float]]]:
    """Get the OpenAI GPT-3 encoder."""
    try:
        from openai import Client as OpenAIClient
    except ImportError:
        raise ImportError(
            "THe default encoder for the EmbeddingDistance class uses the OpenAI API. "
            "Please either install the openai library with `pip install openai` or "
            "provide a custom encoder function (Callable[[str], Sequence[float]])."
        )

    def encode_text(texts: Sequence[str]) -> Sequence[Sequence[float]]:
        client = OpenAIClient()
        response = client.embeddings.create(
            input=list(texts), model="text-embedding-3-small"
        )
        return [d.embedding for d in response.data]

    return encode_text


class EmbeddingConfig(TypedDict, total=False):
    encoder: Callable[[List[str]], Sequence[Sequence[float]]]
    metric: Literal["cosine", "euclidean", "manhattan", "chebyshev", "hamming"]


class EmbeddingDistance:
    def __init__(
        self,
        config: Optional[EmbeddingConfig] = None,
    ):
        config = config or {}
        self.distance = config.get("metric") or "cosine"
        self.encoder = config.get("encoder") or _get_openai_encoder()

    def evaluate(
        self,
        prediction: str,
        reference: str,
    ) -> float:
        try:
            import numpy as np
        except ImportError:
            raise ImportError(
                "The EmbeddingDistance class requires NumPy. Please install it with "
                "`pip install numpy`."
            )
        embeddings = self.encoder([prediction, reference])
        vector = np.array(embeddings)
        return self._compute_distance(vector[0], vector[1]).item()

    def _compute_distance(self, a: np.ndarray, b: np.ndarray) -> np.floating:
        if self.distance == "cosine":
            return self._cosine_distance(a, b)  # type: ignore
        elif self.distance == "euclidean":
            return self._euclidean_distance(a, b)
        elif self.distance == "manhattan":
            return self._manhattan_distance(a, b)
        elif self.distance == "chebyshev":
            return self._chebyshev_distance(a, b)
        elif self.distance == "hamming":
            return self._hamming_distance(a, b)
        else:
            raise ValueError(f"Invalid distance metric: {self.distance}")

    @staticmethod
    def _cosine_distance(a: np.ndarray, b: np.ndarray) -> np.ndarray:
        """Compute the cosine distance between two vectors.

        Args:
            a (np.ndarray): The first vector.
            b (np.ndarray): The second vector.

        Returns:
            np.ndarray: The cosine distance.
        """
        return 1.0 - cosine_similarity([a], [b])

    @staticmethod
    def _euclidean_distance(a: np.ndarray, b: np.ndarray) -> np.floating:
        """Compute the Euclidean distance between two vectors.

        Args:
            a (np.ndarray): The first vector.
            b (np.ndarray): The second vector.

        Returns:
            np.floating: The Euclidean distance.
        """
        return np.linalg.norm(a - b)

    @staticmethod
    def _manhattan_distance(a: np.ndarray, b: np.ndarray) -> np.floating:
        """Compute the Manhattan distance between two vectors.

        Args:
            a (np.ndarray): The first vector.
            b (np.ndarray): The second vector.

        Returns:
            np.floating: The Manhattan distance.
        """
        return np.sum(np.abs(a - b))

    @staticmethod
    def _chebyshev_distance(a: np.ndarray, b: np.ndarray) -> np.floating:
        """Compute the Chebyshev distance between two vectors.

        Args:
            a (np.ndarray): The first vector.
            b (np.ndarray): The second vector.

        Returns:
            np.floating: The Chebyshev distance.
        """
        return np.max(np.abs(a - b))

    @staticmethod
    def _hamming_distance(a: np.ndarray, b: np.ndarray) -> np.floating:
        """Compute the Hamming distance between two vectors.

        Args:
            a (np.ndarray): The first vector.
            b (np.ndarray): The second vector.

        Returns:
            np.floating: The Hamming distance.
        """
        return np.mean(a != b)