"""
Model module for bi-encoder models.

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

from abc import ABC, abstractmethod
from collections.abc import Iterable, Sequence
from dataclasses import dataclass
from string import punctuation
from typing import Literal, Self, overload

import torch
from transformers import BatchEncoding

from ..base import LightningIRModel, LightningIROutput
from ..modeling_utils.batching import _batch_elementwise_scoring
from .bi_encoder_config import BiEncoderConfig, MultiVectorBiEncoderConfig, SingleVectorBiEncoderConfig


@dataclass
class BiEncoderEmbedding:
    """Dataclass containing embeddings and the encoding for 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 = None
    """Mask tensor designating which vectors should be ignored during scoring."""
    encoding: BatchEncoding | None = None
    """Tokenizer encodings used to generate the embeddings."""
    ids: list[str] | None = None
    """list of ids for the embeddings, e.g., query or document ids."""

    @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.

        Args:
            device (torch.device | BiEncoderEmbedding): Device to move the embeddings to or another instance to move to
                the same device.
        Returns:
            Self: The instance with embeddings moved to the specified device.
        """
        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.

        Returns:
            torch.device: Device of the embeddings.
        Raises:
            ValueError: If the embeddings and scoring_mask are not on the same device.
        """
        return self.embeddings.device

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

        Yields:
            tuple[str, torch.Tensor]: tuple of attribute name and its value.
        """
        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."""
    similarity: torch.Tensor | None = None
    """Similarity scores between query and document embeddings."""


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`.

        Args:
            config (BiEncoderConfig): Configuration for the bi-encoder model.
        Raises:
            ValueError: If the similarity function is not supported.
        """
        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.

        Args:
            query_encoding (BatchEncoding | None): Tokenizer encodings for the queries. Defaults to None.
            doc_encoding (BatchEncoding | None): Tokenizer encodings for the documents. Defaults to None.
            num_docs (Sequence[int] | int | None): 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.
        Returns:
            BiEncoderOutput: Output of the model containing query and document embeddings and relevance scores.
        """
        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)
        output = BiEncoderOutput(query_embeddings=query_embeddings, doc_embeddings=doc_embeddings)
        if doc_embeddings is not None and query_embeddings is not None:
            output = self.score(output, num_docs)
        return output

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

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

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

        Args:
            encoding (BatchEncoding): Tokenizer encodings for the documents.
        Returns:
            BiEncoderEmbedding: Document embeddings and scoring mask.
        """
        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.

        Args:
            query_embeddings (BiEncoderEmbedding): Embeddings of the queries.
            doc_embeddings (BiEncoderEmbedding): Embeddings of the documents.
            num_docs (Sequence[int] | int | None): Specifies how many documents are passed per query. If a sequence of
                integers, `len(num_docs)` 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.
        Returns:
            torch.Tensor: Similarity scores between all query and document embedding vector pairs.
        """
        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_elementwise_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_elementwise_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_elementwise_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.

        Args:
            encoding (BatchEncoding): Tokenizer encodings for the text sequence.
            input_type (Literal["query", "doc"]): type of input, either "query" or "doc".
        Returns:
            BiEncoderEmbedding: Embeddings and scoring mask for the text sequence.
        Raises:
            NotImplementedError: If the method is not implemented.
        """
        raise NotImplementedError

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

        Args:
            output (BiEncoderOutput): Output containing embeddings and scoring mask.
            num_docs (Sequence[int] | int | None): 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.
        Returns:
            torch.Tensor: Relevance scores.
        Raises:
            NotImplementedError: If the method is not implemented.
        """
        raise NotImplementedError


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`.

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

    def score(
        self,
        output: BiEncoderOutput,
        num_docs: Sequence[int] | int | None = None,
    ) -> BiEncoderOutput:
        """Compute relevance scores between queries and documents.

        Args:
            output (BiEncoderOutput): Output containing embeddings and scoring mask.
            num_docs (Sequence[int] | int | None): 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.
        Returns:
            BiEncoderOutput: Output containing relevance scores.
        Raises:
            ValueError: If query or document embeddings are not provided in the output.
        """
        if output.query_embeddings is None or output.doc_embeddings is None:
            raise ValueError("Expected query and document embeddings in BiEncoderOutput")
        similarity = self.compute_similarity(output.query_embeddings, output.doc_embeddings, num_docs)
        output.scores = similarity.view(-1)
        output.similarity = similarity
        return output


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`.

        Args:
            config (MultiVectorBiEncoderConfig): Configuration for the multi-vector bi-encoder model
        Raises:
            ValueError: If mask scoring tokens are specified in the configuration but the tokenizer is not available
            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 as err:
                raise ValueError(
                    "Can't use token scoring masking if the checkpoint does not have a tokenizer."
                ) from err
            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,
        aggregation_function: Literal["max", "sum", "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 aggregation_function == "max":
            scores = scores.masked_fill(~mask, float("-inf"))
            return scores.amax(dim, keepdim=True)
        scores = scores.masked_fill(~mask, 0)
        if aggregation_function == "sum":
            return scores.sum(dim, keepdim=True)
        num_non_masked = mask.sum(dim, keepdim=True)
        if aggregation_function == "mean":
            return scores.masked_fill(num_non_masked == 0, 0).sum(dim, keepdim=True) / num_non_masked.clamp(min=1)
        raise ValueError(f"Unknown aggregation {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.

        Args:
            encoding (BatchEncoding): Tokenizer encodings for the text sequence.
            input_type (Literal["query", "doc"]): type of input, either "query" or "doc".
        Returns:
            torch.Tensor: Scoring mask.
        """
        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.

        Args:
            similarity (torch.Tensor): Query--document similarity matrix.
            query_scoring_mask (torch.Tensor): Which query vectors should be masked out during scoring.
            doc_scoring_mask (torch.Tensor): Which document vectors should be masked out during scoring.
        Returns:
            torch.Tensor: Aggregated similarity scores.
        """
        scores = similarity
        if scores.shape[-1] > 1:
            scores = self._aggregate(scores, doc_scoring_mask, self.config.doc_aggregation_function, -1)
        if scores.shape[-2] > 1:
            num_docs_t = self._parse_num_docs(
                query_scoring_mask.shape[0], doc_scoring_mask.shape[0], num_docs, similarity.device
            )
            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, output: BiEncoderOutput, num_docs: Sequence[int] | int | None = None) -> BiEncoderOutput:
        """Compute relevance scores between queries and documents.

        Args:
            output (BiEncoderOutput): Output containing embeddings and scoring mask.
            num_docs (Sequence[int] | int | None): 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.
        Returns:
            BiEncoderOutput: Output containing relevance scores.
        Raises:
            ValueError: If query or document embeddings are not provided in the output.
            ValueError: If scoring masks are not provided for the embeddings.
        """
        query_embeddings = output.query_embeddings
        doc_embeddings = output.doc_embeddings
        if query_embeddings is None or doc_embeddings is None:
            raise ValueError("Embeddings expected for scoring multi-vector embeddings")
        if query_embeddings.scoring_mask is None or doc_embeddings.scoring_mask is None:
            raise ValueError("Scoring masks expected for scoring multi-vector embeddings")

        similarity = self.compute_similarity(query_embeddings, doc_embeddings, num_docs)
        scores = self.aggregate_similarity(
            similarity, query_embeddings.scoring_mask, doc_embeddings.scoring_mask, num_docs
        )
        output.scores = scores
        output.similarity = similarity
        return output
