"""
Model module for bi-encoder models.

This module defines the model class used to implement bi-encoder models.
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass
from functools import wraps
from string import punctuation
from typing import Callable, Iterable, Literal, Self, Sequence, Tuple, Type, overload

import torch
from transformers import BatchEncoding

from ..base import LightningIRModel, LightningIROutput
from .bi_encoder_config import BiEncoderConfig, MultiVectorBiEncoderConfig, SingleVectorBiEncoderConfig


@dataclass
class BiEncoderEmbedding:
    """Dataclass containing embeddings and the encoding for single-vector bi-encoder models."""

    embeddings: torch.Tensor
    """Embedding tensor generated by a bi-encoder model of shape [batch_size x seq_len x hidden_size]. The sequence
    length varies depending on the pooling strategy and the hidden size varies depending on the projection settings."""
    scoring_mask: torch.Tensor | None
    """Mask tensor designating which vectors should be ignored during scoring."""
    encoding: BatchEncoding | None
    """Tokenizer encodings used to generate the embeddings."""

    @overload
    def to(self, device: torch.device, /) -> Self: ...

    @overload
    def to(self, other: Self, /) -> Self: ...

    def to(self, device) -> Self:
        """Moves the embeddings to the specified device.

        :param device: Device to move the embeddings to or another instance to move to the same device
        :type device: torch.device | BiEncoderEmbedding
        :return: Self
        :rtype: BiEncoderEmbedding
        """
        if isinstance(device, BiEncoderEmbedding):
            device = device.device
        self.embeddings = self.embeddings.to(device)
        self.scoring_mask = self.scoring_mask.to(device) if self.scoring_mask is not None else None
        self.encoding = self.encoding.to(device) if self.encoding is not None else None
        return self

    @property
    def device(self) -> torch.device:
        """Returns the device of the embeddings.

        :raises ValueError: If the embeddings and scoring_mask are not on the same device
        :return: The device of the embeddings
        :rtype: torch.device
        """
        return self.embeddings.device

    def items(self) -> Iterable[Tuple[str, torch.Tensor]]:
        """Iterates over the embeddings attributes and their values like `dict.items()`.

        :yield: Tuple of attribute name and its value
        :rtype: Iterator[Iterable[Tuple[str, torch.Tensor]]]
        """
        for field in self.__dataclass_fields__:
            yield field, getattr(self, field)


@dataclass
class BiEncoderOutput(LightningIROutput):
    """Dataclass containing the output of a bi-encoder model."""

    query_embeddings: BiEncoderEmbedding | None = None
    """Query embeddings generated by the model."""
    doc_embeddings: BiEncoderEmbedding | None = None
    """Document embeddings generated by the model."""


def _batch_scoring(
    similarity_function: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
) -> Callable[[torch.Tensor, torch.Tensor], torch.Tensor]:
    """Helper function to batch similarity functions to avoid memory issues with large batch sizes or high numbers
    of documents per query."""
    BATCH_SIZE = 16384

    @wraps(similarity_function)
    def batch_similarity_function(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        if x.shape[0] <= BATCH_SIZE:
            return similarity_function(x, y)
        out = torch.zeros(x.shape[0], x.shape[1], y.shape[2], device=x.device, dtype=x.dtype)
        for i in range(0, x.shape[0], BATCH_SIZE):
            out[i : i + BATCH_SIZE] = similarity_function(x[i : i + BATCH_SIZE], y[i : i + BATCH_SIZE])
        return out

    return batch_similarity_function


class BiEncoderModel(LightningIRModel, ABC):
    """A bi-encoder model that encodes queries and documents separately and computes a relevance score between them.
    See :class:`.BiEncoderConfig` for configuration options."""

    config_class: Type[BiEncoderConfig] = BiEncoderConfig
    """Configuration class for the bi-encoder model."""

    def __init__(self, config: BiEncoderConfig, *args, **kwargs) -> None:
        """Initializes a bi-encoder model given a :class:`.BiEncoderConfig`.

        :param config: Configuration for the bi-encoder model
        :type config: BiEncoderConfig
        """
        super().__init__(config, *args, **kwargs)
        self.config: BiEncoderConfig

        if self.config.similarity_function == "cosine":
            self.similarity_function = self._cosine_similarity
        elif self.config.similarity_function == "l2":
            self.similarity_function = self._l2_similarity
        elif self.config.similarity_function == "dot":
            self.similarity_function = self._dot_similarity
        else:
            raise ValueError(f"Unknown similarity function {self.config.similarity_function}")

    def forward(
        self,
        query_encoding: BatchEncoding | None,
        doc_encoding: BatchEncoding | None,
        num_docs: Sequence[int] | int | None = None,
    ) -> BiEncoderOutput:
        """Embeds queries and/or documents and computes relevance scores between them if both are provided.

        :param query_encoding: Tokenizer encodings for the queries
        :type query_encoding: BatchEncoding | None
        :param doc_encoding: Tokenizer encodings for the documents
        :type doc_encoding: BatchEncoding | None
        :param num_docs: Specifies how many documents are passed per query. If a sequence of integers, `len(num_doc)`
            should be equal to the number of queries and `sum(num_docs)` equal to the number of documents, i.e., the
            sequence contains one value per query specifying the number of documents for that query. If an integer,
            assumes an equal number of documents per query. If None, tries to infer the number of documents by dividing
            the number of documents by the number of queries, defaults to None
        :type num_docs: Sequence[int] | int | None, optional
        :return: Output of the model
        :rtype: BiEncoderOutput
        """
        query_embeddings = None
        if query_encoding is not None:
            query_embeddings = self.encode_query(query_encoding)
        doc_embeddings = None
        if doc_encoding is not None:
            doc_embeddings = self.encode_doc(doc_encoding)
        scores = None
        if doc_embeddings is not None and query_embeddings is not None:
            scores = self.score(query_embeddings, doc_embeddings, num_docs)
        return BiEncoderOutput(scores=scores, query_embeddings=query_embeddings, doc_embeddings=doc_embeddings)

    def encode_query(self, encoding: BatchEncoding) -> BiEncoderEmbedding:
        """Encodes tokenized queries.

        :param encoding: Tokenizer encodings for the queries
        :type encoding: BatchEncoding
        :return: Query embeddings and scoring mask
        :rtype: BiEncoderEmbedding
        """
        return self.encode(encoding=encoding, input_type="query")

    def encode_doc(self, encoding: BatchEncoding) -> BiEncoderEmbedding:
        """Encodes tokenized documents.

        :param encoding: Tokenizer encodings for the documents
        :type encoding: BatchEncoding
        :return: Query embeddings and scoring mask
        :rtype: BiEncoderEmbedding
        """
        return self.encode(encoding=encoding, input_type="doc")

    def _parse_num_docs(
        self, query_shape: int, doc_shape: int, num_docs: int | Sequence[int] | None, device: torch.device | None = None
    ) -> torch.Tensor:
        """Helper function to parse the number of documents per query."""
        if isinstance(num_docs, int):
            num_docs = [num_docs] * query_shape
        if isinstance(num_docs, list):
            if sum(num_docs) != doc_shape or len(num_docs) != query_shape:
                raise ValueError("Num docs does not match doc embeddings")
        if num_docs is None:
            if doc_shape % query_shape != 0:
                raise ValueError("Docs are not evenly distributed in _batch, but no num_docs provided")
            num_docs = [doc_shape // query_shape] * query_shape
        return torch.tensor(num_docs, device=device)

    def compute_similarity(
        self,
        query_embeddings: BiEncoderEmbedding,
        doc_embeddings: BiEncoderEmbedding,
        num_docs: Sequence[int] | int | None = None,
    ) -> torch.Tensor:
        """Computes the similarity score between all query and document embedding vector pairs.

        :param query_embeddings: Embeddings of the queries
        :type query_embeddings: BiEncoderEmbedding
        :param doc_embeddings: Embeddings of the documents
        :type doc_embeddings: BiEncoderEmbedding
        :param num_docs: Specifies how many documents are passed per query. If a sequence of integers, `len(num_doc)`
            should be equal to the number of queries and `sum(num_docs)` equal to the number of documents, i.e., the
            sequence contains one value per query specifying the number of documents for that query. If an integer,
            assumes an equal number of documents per query. If None, tries to infer the number of documents by dividing
            the number of documents by the number of queries, defaults to None
        :type num_docs: Sequence[int] | int | None, optional
        :return: Similarity scores between all query and document embedding vector pairs
        :rtype: torch.Tensor
        """
        num_docs_t = self._parse_num_docs(
            query_embeddings.embeddings.shape[0], doc_embeddings.embeddings.shape[0], num_docs, query_embeddings.device
        )
        query_emb = query_embeddings.embeddings.repeat_interleave(num_docs_t, dim=0).unsqueeze(2)
        doc_emb = doc_embeddings.embeddings.unsqueeze(1)
        similarity = self.similarity_function(query_emb, doc_emb)
        return similarity

    @staticmethod
    @_batch_scoring
    @torch.autocast(device_type="cuda", enabled=False)
    def _cosine_similarity(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        return torch.nn.functional.cosine_similarity(x, y, dim=-1)

    @staticmethod
    @_batch_scoring
    @torch.autocast(device_type="cuda", enabled=False)
    def _l2_similarity(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        return -1 * torch.cdist(x, y).squeeze(-2)

    @staticmethod
    @_batch_scoring
    @torch.autocast(device_type="cuda", enabled=False)
    def _dot_similarity(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        return torch.matmul(x, y.transpose(-1, -2)).squeeze(-2)

    @abstractmethod
    def encode(self, encoding: BatchEncoding, input_type: Literal["query", "doc"]) -> BiEncoderEmbedding:
        """Encodes a batched tokenized text sequences and returns the embeddings and scoring mask.

        :param encoding: Tokenizer encodings for the text sequence
        :type encoding: BatchEncoding
        :param input_type: Type of input, either "query" or "doc"
        :type input_type: Literal["query", "doc"]
        :return: Embeddings and scoring mask
        :rtype: BiEncoderEmbedding
        """
        pass

    @abstractmethod
    def score(
        self,
        query_embeddings: BiEncoderEmbedding,
        doc_embeddings: BiEncoderEmbedding,
        num_docs: Sequence[int] | int | None = None,
    ) -> torch.Tensor:
        """Compute relevance scores between queries and documents.

        :param query_embeddings: Embeddings and scoring mask for the queries
        :type query_embeddings: BiEncoderEmbedding
        :param doc_embeddings: Embeddings and scoring mask for the documents
        :type doc_embeddings: BiEncoderEmbedding
        :param num_docs: Specifies how many documents are passed per query. If a sequence of integers, `len(num_doc)`
            should be equal to the number of queries and `sum(num_docs)` equal to the number of documents, i.e., the
            sequence contains one value per query specifying the number of documents for that query. If an integer,
            assumes an equal number of documents per query. If None, tries to infer the number of documents by dividing
            the number of documents by the number of queries, defaults to None
        :type num_docs: Sequence[int] | int | None, optional
        :return: Relevance scores
        :rtype: torch.Tensor
        """
        pass


class SingleVectorBiEncoderModel(BiEncoderModel):
    """A bi-encoder model that encodes queries and documents separately, pools the contextualized embeddings into a
    single vector, and computes a relevance score based on the similarities between the two vectors. See
    :class:`.SingleVectorBiEncoderConfig` for configuration options."""

    config_class: Type[SingleVectorBiEncoderConfig] = SingleVectorBiEncoderConfig
    """Configuration class for the single-vector bi-encoder model."""

    def __init__(self, config: SingleVectorBiEncoderConfig, *args, **kwargs) -> None:
        """Initializes a single-vector bi-encoder model given a :class:`.SingleVectorBiEncoderConfig`.

        :param config: Configuration for the single-vector bi-encoder model
        :type config: SingleVectorBiEncoderConfig
        """
        super().__init__(config, *args, **kwargs)

    def score(
        self,
        query_embeddings: BiEncoderEmbedding,
        doc_embeddings: BiEncoderEmbedding,
        num_docs: Sequence[int] | int | None = None,
    ) -> torch.Tensor:
        """Compute relevance scores between queries and documents.

        :param query_embeddings: Embeddings and scoring mask for the queries
        :type query_embeddings: BiEncoderEmbedding
        :param doc_embeddings: Embeddings and scoring mask for the documents
        :type doc_embeddings: BiEncoderEmbedding
        :param num_docs: Specifies how many documents are passed per query. If a sequence of integers, `len(num_doc)`
            should be equal to the number of queries and `sum(num_docs)` equal to the number of documents, i.e., the
            sequence contains one value per query specifying the number of documents for that query. If an integer,
            assumes an equal number of documents per query. If None, tries to infer the number of documents by dividing
            the number of documents by the number of queries, defaults to None
        :type num_docs: Sequence[int] | int | None, optional
        :return: Relevance scores
        :rtype: torch.Tensor
        """
        return self.compute_similarity(query_embeddings, doc_embeddings, num_docs).view(-1)


class MultiVectorBiEncoderModel(BiEncoderModel):

    config_class: Type[MultiVectorBiEncoderConfig] = MultiVectorBiEncoderConfig
    """Configuration class for the single-vector bi-encoder model."""

    supports_retrieval_models = []

    def __init__(self, config: MultiVectorBiEncoderConfig, *args, **kwargs) -> None:
        """Initializes a multi-vector bi-encoder model given a :class:`.MultiVectorBiEncoderConfig`.

        :param config: Configuration for the multi-vector bi-encoder model
        :type config: MultiVectorBiEncoderConfig
        :raises ValueError: If mask scoring tokens are specified in the configuration but the tokenizer is not available
        :raises ValueError: If the specified mask scoring tokens are not in the tokenizer vocab
        """
        super().__init__(config, *args, **kwargs)

        self.query_mask_scoring_input_ids: torch.Tensor | None = None
        self.doc_mask_scoring_input_ids: torch.Tensor | None = None

        # Adds the mask scoring input ids to the model if they are specified in the configuration.
        for sequence in ("query", "doc"):
            mask_scoring_tokens = getattr(self.config, f"{sequence}_mask_scoring_tokens")
            if mask_scoring_tokens is None:
                continue
            if mask_scoring_tokens == "punctuation":
                mask_scoring_tokens = list(punctuation)
            try:
                from transformers import AutoTokenizer

                tokenizer = AutoTokenizer.from_pretrained(self.config.name_or_path)
            except OSError:
                raise ValueError("Can't use token scoring masking if the checkpoint does not have a tokenizer.")
            mask_scoring_input_ids = []
            for token in mask_scoring_tokens:
                if token not in tokenizer.vocab:
                    raise ValueError(f"Token {token} not in tokenizer vocab")
                mask_scoring_input_ids.append(tokenizer.vocab[token])
            setattr(
                self,
                f"{sequence}_mask_scoring_input_ids",
                torch.tensor(mask_scoring_input_ids, dtype=torch.long),
            )

    def _expand_mask(self, shape: torch.Size, mask: torch.Tensor, dim: int) -> torch.Tensor:
        """Helper function to expand the mask to the shape of the similarity scores."""
        if mask.ndim == len(shape):
            return mask
        if mask.ndim > len(shape):
            raise ValueError("Mask has too many dimensions")
        fill_values = len(shape) - mask.ndim + 1
        new_shape = [*mask.shape[:-1]] + [1] * fill_values
        new_shape[dim] = mask.shape[-1]
        return mask.view(*new_shape)

    def _aggregate(
        self,
        scores: torch.Tensor,
        mask: torch.Tensor,
        query_aggregation_function: Literal["max", "sum", "mean", "harmonic_mean"],
        dim: int,
    ) -> torch.Tensor:
        """Helper function to aggregate similarity scores over query and document embeddings."""
        mask = self._expand_mask(scores.shape, mask, dim)
        if query_aggregation_function == "max":
            scores = scores.masked_fill(~mask, float("-inf"))
            return scores.amax(dim, keepdim=True)
        if query_aggregation_function == "sum":
            scores = scores.masked_fill(~mask, 0)
            return scores.sum(dim, keepdim=True)
        num_non_masked = mask.sum(dim, keepdim=True)
        if query_aggregation_function == "mean":
            return torch.where(num_non_masked == 0, 0, scores.sum(dim, keepdim=True) / num_non_masked)
        if query_aggregation_function == "harmonic_mean":
            return torch.where(
                num_non_masked == 0,
                0,
                num_non_masked / (1 / scores).sum(dim, keepdim=True),
            )
        raise ValueError(f"Unknown aggregation {query_aggregation_function}")

    def scoring_mask(self, encoding: BatchEncoding, input_type: Literal["query", "doc"]) -> torch.Tensor:
        """Computes a scoring mask for batched tokenized text sequences which is used in the scoring function to mask
        out vectors during scoring.

        :param encoding: Tokenizer encodings for the text sequence
        :type encoding: BatchEncoding
        :param input_type: Type of input, either "query" or "doc"
        :type input_type: Literal["query", "doc"]
        :return: Scoring mask
        :rtype: torch.Tensor
        """
        input_ids = encoding["input_ids"]
        attention_mask = encoding["attention_mask"]
        scoring_mask = attention_mask
        if scoring_mask is None:
            scoring_mask = torch.ones_like(input_ids, dtype=torch.bool)
        scoring_mask = scoring_mask.bool()
        mask_scoring_input_ids = getattr(self, f"{input_type}_mask_scoring_input_ids")
        if mask_scoring_input_ids is not None:
            ignore_mask = input_ids[..., None].eq(mask_scoring_input_ids.to(input_ids.device)).any(-1)
            scoring_mask = scoring_mask & ~ignore_mask
        return scoring_mask

    def aggregate_similarity(
        self,
        similarity: torch.Tensor,
        query_scoring_mask: torch.Tensor,
        doc_scoring_mask: torch.Tensor,
        num_docs: int | Sequence[int] | None = None,
    ) -> torch.Tensor:
        """Aggregates the matrix of query-document similarities into a single score based on the configured aggregation
        strategy.

        :param similarity: Query-document similarity matrix
        :type similarity: torch.Tensor
        :param query_scoring_mask: Which query vectors should be masked out during scoring
        :type query_scoring_mask: torch.Tensor
        :param doc_scoring_mask: Which doucment vectors should be masked out during scoring
        :type doc_scoring_mask: torch.Tensor
        :return: Aggregated similarity scores
        :rtype: torch.Tensor
        """
        num_docs_t = self._parse_num_docs(
            query_scoring_mask.shape[0], doc_scoring_mask.shape[0], num_docs, similarity.device
        )
        scores = self._aggregate(similarity, doc_scoring_mask, self.config.doc_aggregation_function, -1)
        repeated_query_scoring_mask = query_scoring_mask.repeat_interleave(num_docs_t, dim=0)
        scores = self._aggregate(scores, repeated_query_scoring_mask, self.config.query_aggregation_function, -2)
        return scores.view(scores.shape[0])

    def score(
        self,
        query_embeddings: BiEncoderEmbedding,
        doc_embeddings: BiEncoderEmbedding,
        num_docs: Sequence[int] | int | None = None,
    ) -> torch.Tensor:
        """Compute relevance scores between queries and documents.

        :param query_embeddings: Embeddings and scoring mask for the queries
        :type query_embeddings: BiEncoderEmbedding
        :param doc_embeddings: Embeddings and scoring mask for the documents
        :type doc_embeddings: BiEncoderEmbedding
        :param num_docs: Specifies how many documents are passed per query. If a sequence of integers, `len(num_doc)`
            should be equal to the number of queries and `sum(num_docs)` equal to the number of documents, i.e., the
            sequence contains one value per query specifying the number of documents for that query. If an integer,
            assumes an equal number of documents per query. If None, tries to infer the number of documents by dividing
            the number of documents by the number of queries, defaults to None
        :type num_docs: Sequence[int] | int | None, optional
        :return: Relevance scores
        :rtype: torch.Tensor
        """
        similarity = self.compute_similarity(query_embeddings, doc_embeddings, num_docs)
        if query_embeddings.scoring_mask is None or doc_embeddings.scoring_mask is None:
            raise ValueError("Scoring masks expected for scoring multi-vector embeddings")
        return self.aggregate_similarity(similarity, query_embeddings.scoring_mask, doc_embeddings.scoring_mask)
