"""Configuration, model, and embedding for COIL (Contextualized Inverted list) type models. Originally proposed in
`COIL: Revisit Exact Lexical Match in Information Retrieval with Contextualized Inverted list \
<https://arxiv.org/abs/2104.07186>`_."""

from collections.abc import Sequence
from dataclasses import dataclass
from typing import Literal

import torch
from transformers import BatchEncoding

from ...bi_encoder import (
    BiEncoderEmbedding,
    BiEncoderOutput,
    MultiVectorBiEncoderConfig,
    MultiVectorBiEncoderModel,
    SingleVectorBiEncoderConfig,
    SingleVectorBiEncoderModel,
)


@dataclass
class CoilEmbedding(BiEncoderEmbedding):
    """Dataclass containing embeddings and the encoding for COIL models."""

    embeddings: torch.Tensor
    """Raw embeddings of the COIL model. Should not be used directly for scoring."""
    token_embeddings: torch.Tensor | None = None
    """Token embeddings of the COIL model."""
    cls_embeddings: torch.Tensor | None = None
    """Separate [CLS] token embeddings."""


@dataclass
class CoilOutput(BiEncoderOutput):
    """Dataclass containing the output of a COIL model."""

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


class CoilConfig(MultiVectorBiEncoderConfig):
    """Configuration class for COIL models."""

    model_type = "coil"
    """Model type for COIL models."""

    def __init__(
        self,
        query_length: int | None = 32,
        doc_length: int | None = 512,
        similarity_function: Literal["cosine", "dot"] = "dot",
        add_marker_tokens: bool = False,
        token_embedding_dim: int = 32,
        cls_embedding_dim: int = 768,
        projection: Literal["linear", "linear_no_bias"] = "linear",
        **kwargs,
    ) -> None:
        """A COIL model encodes queries and documents separately, and computes a similarity score using the maximum
        similarity ...

        Args:
            query_length (int | None): Maximum number of tokens per query. If None does not truncate. Defaults to 32.
            doc_length (int | None): Maximum number of tokens per document. If None does not truncate. Defaults to 512.
            similarity_function (Literal["cosine", "dot"]): Similarity function to compute scores between query and
                document embeddings. Defaults to "dot".
            add_marker_tokens (bool): Whether to add extra marker tokens [Q] / [D] to queries / documents.
                Defaults to False.
            token_embedding_dim (int | None): The output embedding dimension for tokens. Defaults to 32.
            cls_embedding_dim (int | None): The output embedding dimension for the [CLS] token. Defaults to 768.
            projection (Literal["linear", "linear_no_bias"], optional): Whether and how to project the embeddings.
                Defaults to "linear".
        """
        super().__init__(
            query_length=query_length,
            doc_length=doc_length,
            similarity_function=similarity_function,
            add_marker_tokens=add_marker_tokens,
            **kwargs,
        )
        self.projection = projection
        self.token_embedding_dim = token_embedding_dim
        self.cls_embedding_dim = cls_embedding_dim


class CoilModel(MultiVectorBiEncoderModel):
    """Multi-vector COIL model. See :class:`.CoilConfig` for configuration options."""

    config_class = CoilConfig
    """Configuration class for COIL models."""

    def __init__(self, config: CoilConfig, *args, **kwargs) -> None:
        """Initializes a COIL model given a :class:`.CoilConfig` configuration.

        Args:
            config (CoilConfig): Configuration for the COIL model.
        """
        super().__init__(config, *args, **kwargs)
        self.config: CoilConfig
        self.token_projection = torch.nn.Linear(
            self.config.hidden_size,
            self.config.token_embedding_dim,
            bias="no_bias" not in self.config.projection,
        )
        self.cls_projection = torch.nn.Linear(
            self.config.hidden_size,
            self.config.cls_embedding_dim,
            bias="no_bias" not in self.config.projection,
        )

    def encode(self, encoding: BatchEncoding, input_type: Literal["query", "doc"]) -> CoilEmbedding:
        """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.
        """
        embeddings = self._backbone_forward(**encoding).last_hidden_state

        cls_embeddings = self.cls_projection(embeddings[:, [0]])
        token_embeddings = self.token_projection(embeddings[:, 1:])

        scoring_mask = self.scoring_mask(encoding, input_type)
        return CoilEmbedding(
            embeddings, scoring_mask, encoding, cls_embeddings=cls_embeddings, token_embeddings=token_embeddings
        )

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

        Args:
            query_embeddings (CoilEmbedding): CLS embeddings, token embeddings, and scoring mask for the queries.
            doc_embeddings (CoilEmbedding): CLS embeddings, token embeddings, and scoring mask for the documents.
            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."""
        query_embeddings = output.query_embeddings
        doc_embeddings = output.doc_embeddings
        if query_embeddings is None or doc_embeddings is None:
            raise ValueError("Query and document embeddings must be provided for scoring.")
        if query_embeddings.scoring_mask is None or doc_embeddings.scoring_mask is None:
            raise ValueError("Scoring masks expected for scoring multi-vector embeddings")
        if (
            query_embeddings.cls_embeddings is None
            or doc_embeddings.cls_embeddings is None
            or query_embeddings.token_embeddings is None
            or doc_embeddings.token_embeddings is None
        ):
            raise ValueError("COIL embeddings must contain cls_embeddings and token_embeddings")

        cls_scores = self.compute_similarity(
            BiEncoderEmbedding(query_embeddings.cls_embeddings),
            BiEncoderEmbedding(doc_embeddings.cls_embeddings),
            num_docs,
        ).view(-1)

        token_similarity = self.compute_similarity(
            BiEncoderEmbedding(query_embeddings.token_embeddings),
            BiEncoderEmbedding(doc_embeddings.token_embeddings),
            num_docs,
        )
        num_docs_t = self._parse_num_docs(
            query_embeddings.embeddings.shape[0], doc_embeddings.embeddings.shape[0], num_docs, query_embeddings.device
        )
        query = query_embeddings.encoding.input_ids.repeat_interleave(num_docs_t, 0)[:, 1:]
        docs = doc_embeddings.encoding.input_ids[:, 1:]
        mask = (query[:, :, None] == docs[:, None, :]).to(token_similarity)
        token_similarity = token_similarity * mask
        token_scores = self.aggregate_similarity(
            token_similarity, query_embeddings.scoring_mask[:, 1:], doc_embeddings.scoring_mask[:, 1:], num_docs
        )

        output.scores = cls_scores + token_scores
        return output


class UniCoilConfig(SingleVectorBiEncoderConfig):
    """Configuration class for UniCOIL models."""

    model_type = "unicoil"
    """Model type for UniCOIL models."""

    def __init__(
        self,
        query_length: int | None = 32,
        doc_length: int | None = 512,
        similarity_function: Literal["cosine", "dot"] = "dot",
        projection: Literal["linear", "linear_no_bias"] = "linear",
        **kwargs,
    ) -> None:
        """A UniCOIL model encodes queries and documents separately, and computes a similarity score using the maximum
        similarity of token embeddings between query and document.

        Args:
            query_length (int | None): Maximum query length in number of tokens. Defaults to 32.
            doc_length (int | None): Maximum document length in number of tokens. Defaults to 512.
            similarity_function (Literal["cosine", "dot"]): Similarity function to compute scores between query and
                document embeddings. Defaults to "dot".
            projection (Literal["linear", "linear_no_bias"], optional): Whether and how to project the embeddings.
                Defaults to "linear".
        """
        super().__init__(
            query_length=query_length,
            doc_length=doc_length,
            similarity_function=similarity_function,
            **kwargs,
        )
        self.projection = projection

    @property
    def embedding_dim(self) -> int:
        vocab_size = getattr(self, "vocab_size", None)
        if vocab_size is None:
            raise ValueError("Unable to determine embedding dimension.")
        return vocab_size

    @embedding_dim.setter
    def embedding_dim(self, value: int) -> None:
        pass


class UniCoilModel(SingleVectorBiEncoderModel):
    """Single-vector UniCOIL model. See :class:`.UniCoilConfig` for configuration options."""

    config_class = UniCoilConfig
    """Configuration class for UniCOIL models."""

    def __init__(self, config: UniCoilConfig, *args, **kwargs) -> None:
        """Initializes a UniCOIL model given a :class:`.UniCoilConfig` configuration.

        Args:
            config (UniCoilConfig): Configuration for the UniCOIL model.
        """
        super().__init__(config, *args, **kwargs)
        self.config: UniCoilConfig
        self.token_projection = torch.nn.Linear(
            self.config.hidden_size, 1, bias="no_bias" not in self.config.projection
        )

    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.
        """
        contextualized_embeddings = self._backbone_forward(**encoding).last_hidden_state

        token_weights = self.token_projection(contextualized_embeddings).squeeze(-1)
        if encoding["attention_mask"] is not None:
            token_weights = token_weights.masked_fill(~(encoding["attention_mask"].bool()), 0)
        token_weights = torch.relu(token_weights)
        embeddings = torch.zeros(encoding.input_ids.shape[0], self.config.vocab_size, device=token_weights.device)
        embeddings = embeddings.scatter(1, encoding.input_ids, token_weights)
        return BiEncoderEmbedding(embeddings[:, None], None, encoding)
