import torch
from torch import nn
from transformers.models.paligemma.modeling_paligemma import PaliGemmaForConditionalGeneration, PaliGemmaPreTrainedModel


class BiPaliLast(PaliGemmaPreTrainedModel):
    def __init__(self, config):
        super(BiPaliLast, self).__init__(config=config)
        self.model: PaliGemmaForConditionalGeneration = PaliGemmaForConditionalGeneration(config)
        self.pooling_strategy = "last"
        self.main_input_name = "doc_input_ids"

    def forward(self, *args, **kwargs):
        """
        Forward pass through Llama and the linear layer for dimensionality reduction

        Args:
        - input_ids (torch.LongTensor): The input tokens tensor.
        - attention_mask (torch.LongTensor): The attention mask tensor.

        Returns:
        - torch.Tensor: Embeddings of shape (batch_size, num_tokens, dim)
        """
        outputs = self.model(*args, output_hidden_states=True, **kwargs)
        last_hidden_states = outputs.hidden_states[-1]  # (batch_size, sequence_length, hidden_size)
        # pooling - last token
        proj = last_hidden_states[:, -1, :]
        # normalize l2 norm
        proj = proj / proj.norm(dim=-1, keepdim=True)
        return proj


class BiPaliMean(PaliGemmaPreTrainedModel):
    def __init__(self, config):
        super(BiPaliMean, self).__init__(config=config)
        self.model: PaliGemmaForConditionalGeneration = PaliGemmaForConditionalGeneration(config)
        self.pooling_strategy = "mean"
        self.main_input_name = "doc_input_ids"

    def forward(self, *args, **kwargs):
        """
        Forward pass through Llama and the linear layer for dimensionality reduction

        Args:
        - input_ids (torch.LongTensor): The input tokens tensor.
        - attention_mask (torch.LongTensor): The attention mask tensor.

        Returns:
        - torch.Tensor: Embeddings of shape (batch_size, num_tokens, dim)
        """
        outputs = self.model(*args, output_hidden_states=True, **kwargs)
        last_hidden_states = outputs.hidden_states[-1]  # (batch_size, sequence_length, hidden_size)
        # pooling -mean on attention mask==1
        proj = torch.sum(last_hidden_states * kwargs["attention_mask"].unsqueeze(-1), dim=1) / torch.sum(
            kwargs["attention_mask"], dim=1, keepdim=True
        )
        proj = proj / proj.norm(dim=-1, keepdim=True)
        return proj


class ColPali(PaliGemmaPreTrainedModel):
    def __init__(self, config):
        super(ColPali, self).__init__(config=config)
        self.model: PaliGemmaForConditionalGeneration = PaliGemmaForConditionalGeneration(config)
        self.dim = 128
        self.custom_text_proj = nn.Linear(self.model.config.text_config.hidden_size, self.dim)
        self.main_input_name = "doc_input_ids"

    def forward(self, *args, **kwargs):
        """
        Forward pass through Llama and the linear layer for dimensionality reduction

        Args:
        - input_ids (torch.LongTensor): The input tokens tensor.
        - attention_mask (torch.LongTensor): The attention mask tensor.

        Returns:
        - torch.Tensor: Embeddings of shape (batch_size, num_tokens, dim)
        """
        outputs = self.model(*args, output_hidden_states=True, **kwargs)
        last_hidden_states = outputs.hidden_states[-1]  # (batch_size, sequence_length, hidden_size)
        proj = self.custom_text_proj(last_hidden_states)
        # normalize l2 norm
        proj = proj / proj.norm(dim=-1, keepdim=True)
        proj = proj * kwargs["attention_mask"].unsqueeze(-1)
        return proj


class ColNewSiglip(PaliGemmaPreTrainedModel):
    def __init__(self, config):
        super(ColNewSiglip, self).__init__(config=config)
        self.model: PaliGemmaForConditionalGeneration = PaliGemmaForConditionalGeneration(config)
        self.dim = 128
        self.custom_image_proj = nn.Linear(self.model.config.vision_config.projection_dim, self.dim)
        self.custom_text_proj = nn.Linear(self.model.config.text_config.hidden_size, self.dim)
        self.main_input_name = "doc_input_ids"

    def forward(self, *args, **kwargs):
        """
        Forward pass through Llama and the linear layer for dimensionality reduction

        Args:
        - input_ids (torch.LongTensor): The input tokens tensor.
        - attention_mask (torch.LongTensor): The attention mask tensor.

        Returns:
        - torch.Tensor: Embeddings of shape (batch_size, num_tokens, dim)
        """
        # outputs = self.model(*args, output_hidden_states=True, **kwargs)
        if "pixel_values" in kwargs:
            image_features = self.vision_model_output(*args, **kwargs)
            # print(f"Doc: {image_features.shape}")
            proj = self.custom_image_proj(image_features)
            # print(f"Doc proj: {proj.shape}")
            proj = proj / proj.norm(dim=-1, keepdim=True)
        else:
            outputs = self.model(*args, output_hidden_states=True, **kwargs)
            last_hidden_states = outputs.hidden_states[-1]  # (batch_size, sequence_length, hidden_size)
            # print(f"Query: {last_hidden_states.shape}")
            proj = self.custom_text_proj(last_hidden_states)
            # print(f"Query proj: {proj.shape}")
            # normalize l2 norm
            proj = proj / proj.norm(dim=-1, keepdim=True)
            proj = proj * kwargs["attention_mask"].unsqueeze(-1)
        return proj

    def vision_model_output(self, input_ids: torch.LongTensor = None, pixel_values: torch.FloatTensor = None, **kwargs):

        inputs_embeds = self.model.get_input_embeddings()(input_ids)
        # 2. Merge text and images
        if pixel_values is not None and input_ids.shape[1] != 1:
            image_outputs = self.model.vision_tower(pixel_values.to(inputs_embeds.dtype))
            selected_image_feature = image_outputs.last_hidden_state
            image_features = self.model.multi_modal_projector(selected_image_feature)

            return image_features

        raise ValueError("pixel_values is None or input_ids.shape[1] == 1")


class BiNewSiglip(PaliGemmaPreTrainedModel):
    def __init__(self, config):
        super(BiNewSiglip, self).__init__(config=config)
        self.model: PaliGemmaForConditionalGeneration = PaliGemmaForConditionalGeneration(config)
        self.main_input_name = "doc_input_ids"

    def forward(self, *args, **kwargs):
        """
        Forward pass through Llama and the linear layer for dimensionality reduction

        Args:
        - input_ids (torch.LongTensor): The input tokens tensor.
        - attention_mask (torch.LongTensor): The attention mask tensor.

        Returns:
        - torch.Tensor: Embeddings of shape (batch_size, num_tokens, dim)
        """
        # outputs = self.model(*args, output_hidden_states=True, **kwargs)
        if "pixel_values" in kwargs:
            image_features = self.vision_model_output(*args, **kwargs)
            # print(f"Doc: {image_features.shape}")
            # pool image features
            proj = torch.mean(image_features, dim=1)
            # print(f"Doc proj: {proj.shape}")
            norm = proj.norm(dim=-1, keepdim=True)
            proj = proj / norm
        else:
            outputs = self.model(*args, output_hidden_states=True, **kwargs)
            last_hidden_states = outputs.hidden_states[-1]  # (batch_size, sequence_length, hidden_size)
            # pooling -mean on attention mask==1

            proj = torch.sum(last_hidden_states * kwargs["attention_mask"].unsqueeze(-1), dim=1) / torch.sum(
                kwargs["attention_mask"], dim=1, keepdim=True
            )
            # print(f"Query proj: {proj.shape}")
            norm = proj.norm(dim=-1, keepdim=True)
            proj = proj / norm
        return proj

    def vision_model_output(self, input_ids: torch.LongTensor = None, pixel_values: torch.FloatTensor = None, **kwargs):

        inputs_embeds = self.model.get_input_embeddings()(input_ids)
        # 2. Merge text and images
        if pixel_values is not None and input_ids.shape[1] != 1:
            image_outputs = self.model.vision_tower(pixel_values.to(inputs_embeds.dtype))
            selected_image_feature = image_outputs.last_hidden_state
            image_features = self.model.multi_modal_projector(selected_image_feature)

            return image_features

        raise ValueError("pixel_values is None or input_ids.shape[1] == 1")