from functools import partial
import math
from typing import Hashable, List, Optional
import torch
from torch import nn
from torch.nn import functional as F
from transformers import (
    AutoModel,
    AutoTokenizer,
    AutoImageProcessor,
    PreTrainedModel,
    CLIPTextModelWithProjection,
    CLIPVisionModelWithProjection,
    SiglipTextModel,
    SiglipVisionModel
)
from dataclasses import dataclass
from src.models.contrastive_loss.modeling_contrastive_loss import (
    ContrastiveLossModel,
    ContrastiveLossModelOutput
)
from src.custom_utils import get_logger
from src.models.ret2.configuration_ret2 import Ret2Config

logger = get_logger(__name__)


MODULES_WITH_GRAD_CHECKPOINTING = (
    'CLIPTextModel',
    'CLIPVisionModel',
    'SiglipVisionModel',
    'SiglipTextModel',
    'BertModel'
)


@dataclass
class Ret2ModelOutput(ContrastiveLossModelOutput):
    query_feats: Optional[torch.Tensor] = None
    passage_feats: Optional[torch.Tensor] = None


class Ret2Cell(nn.Module):

    def init_gates(self):
        # see https://github.com/google-research/meliad/blob/e8af0543441222c1c4c60d58803511f7cf92908b/transformer/nn_components.py#L230
        std = math.sqrt(0.1 / self.config.hidden_size)
        nn.init.trunc_normal_(self.text_input_gate_w.weight,
                              mean=0.0, std=std, a=-1, b=1)
        nn.init.trunc_normal_(self.vision_input_gate_w.weight,
                              mean=0.0, std=std, a=-1, b=1)
        nn.init.trunc_normal_(self.text_forget_gate_w.weight,
                              mean=0.0, std=std, a=-1, b=1)
        nn.init.trunc_normal_(
            self.vision_forget_gate_w.weight, mean=0.0, std=std, a=-1, b=1)

        nn.init.normal_(self.text_input_gate_w.bias, mean=0.0, std=0.1)
        nn.init.normal_(self.vision_input_gate_w.bias, mean=0.0, std=0.1)
        nn.init.normal_(self.text_forget_gate_w.bias, mean=0.0, std=0.1)
        nn.init.normal_(self.vision_forget_gate_w.bias, mean=0.0, std=0.1)

    def __init__(self, config: Ret2Config):
        super().__init__()
        self.config = config

        self.hidden_state = nn.Parameter(
            torch.FloatTensor(1, 1, config.hidden_size))
        nn.init.normal_(self.hidden_state, mean=0.0, std=0.1)

        self.text_xa = nn.MultiheadAttention(embed_dim=config.hidden_size, kdim=config.text_config.hidden_size, vdim=config.text_config.hidden_size,
                                             num_heads=config.hidden_size // 64, dropout=config.attention_dropout, batch_first=True, bias=False)
        self.vision_xa = nn.MultiheadAttention(embed_dim=config.hidden_size, kdim=config.vision_config.hidden_size, vdim=config.vision_config.hidden_size,
                                               num_heads=config.hidden_size // 64, dropout=config.attention_dropout, batch_first=True, bias=False)

        self.mlp = nn.Sequential(
            nn.Linear(config.hidden_size, config.hidden_size * 4),
            nn.Dropout(config.dropout_p),
            nn.GELU(),
            nn.Linear(config.hidden_size * 4, config.hidden_size),
            nn.Dropout(config.dropout_p)
        )

        self.text_input_gate_w = nn.Linear(
            config.hidden_size, config.hidden_size)
        self.vision_input_gate_w = nn.Linear(
            config.hidden_size, config.hidden_size)

        self.text_forget_gate_w = nn.Linear(
            config.hidden_size, config.hidden_size)
        self.vision_forget_gate_w = nn.Linear(
            config.hidden_size, config.hidden_size)

        self.xa_ln = nn.LayerNorm(config.hidden_size)
        self.mlp_ln = nn.LayerNorm(config.hidden_size)
        self.text_ln = nn.LayerNorm(config.text_config.hidden_size)
        self.vision_ln = nn.LayerNorm(config.vision_config.hidden_size)

        self.register_buffer(
            'input_gate_bias', torch.tensor(config.input_gate_bias))
        self.register_buffer('forget_gate_bias',
                             torch.tensor(config.forget_gate_bias))

        self.init_gates()

    def forward(
        self,
        hidden_state: Optional[torch.Tensor] = None,
        text_key_value: Optional[torch.Tensor] = None,
        text_attention_mask: Optional[torch.Tensor] = None,
        text_mask: Optional[torch.Tensor] = None,
        vision_key_value: Optional[torch.Tensor] = None,
        vision_attention_mask: Optional[torch.Tensor] = None,
        vision_mask: Optional[torch.Tensor] = None
    ):
        if hidden_state is None:
            hidden_state = self.hidden_state.expand(text_key_value.size(0), -1, -1)
            residual = torch.zeros_like(hidden_state)
        else:
            residual = hidden_state

        hidden_state = self.xa_ln(hidden_state)

        text_key_value = self.text_ln(text_key_value)
        text_hidden_state, _ = self.text_xa(
            hidden_state, text_key_value, text_key_value, attn_mask=text_attention_mask)

        vision_key_value = self.vision_ln(vision_key_value)
        vision_hidden_state, _ = self.vision_xa(
            hidden_state, vision_key_value, vision_key_value)

        text_input_gate = F.sigmoid(self.text_input_gate_w(
            text_hidden_state) + self.input_gate_bias) * text_mask
        text_forget_gate = self.text_forget_gate_w(text_hidden_state) * text_mask

        vision_input_gate = F.sigmoid(self.vision_input_gate_w(
            vision_hidden_state) + self.input_gate_bias) * vision_mask
        vision_forget_gate = self.vision_forget_gate_w(vision_hidden_state) * vision_mask

        forget_gate = F.sigmoid(
            text_forget_gate + vision_forget_gate + self.forget_gate_bias)

        hidden_state = residual * forget_gate +\
            text_hidden_state * text_input_gate +\
            vision_hidden_state * vision_input_gate

        return hidden_state + self.mlp(self.mlp_ln(hidden_state))


def get_additive_attn_mask(binary_attn_mask, dtype, force_attn_on_one_item=False):
    if force_attn_on_one_item:
        # force attention on one item to avoid NaN in the attention output
        force_attn_mask = (binary_attn_mask == 0).sum(-1) == binary_attn_mask.size(1)
        binary_attn_mask[force_attn_mask, 0] = 1
    ret = torch.where(binary_attn_mask.bool(), 0, torch.finfo(dtype).min).to(dtype)
    return ret


class Ret2Model(PreTrainedModel):
    config_class = Ret2Config
    supports_gradient_checkpointing = True

    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
        """Enable gradient checkpointing"""
        self.apply(partial(self._set_gradient_checkpointing, value=True))

    def gradient_checkpointing_disable(self):
        """Disable gradient checkpointing"""
        self.apply(partial(self._set_gradient_checkpointing, value=False))

    def _set_gradient_checkpointing(self, module, value=False):
        if value and module.__class__.__name__ in MODULES_WITH_GRAD_CHECKPOINTING:
            module.gradient_checkpointing_enable()

    def init_pretrained_models(self):
        if 'siglip' in self.config.text_name_or_path.lower():
            auto_cls = SiglipTextModel
        elif 'clip-vit-' in self.config.text_name_or_path.lower():
            auto_cls = CLIPTextModelWithProjection
        else:
            auto_cls = AutoModel
        self.text_model = auto_cls.from_pretrained(
            self.config.text_name_or_path, config=self.config.text_config)
        logger.info(
            f"Loaded `{self.text_model.__class__.__name__}` from {self.config.text_name_or_path}")

        if 'siglip' in self.config.vision_name_or_path.lower():
            auto_cls = SiglipVisionModel
        elif 'clip-vit-' in self.config.vision_name_or_path.lower():
            auto_cls = CLIPVisionModelWithProjection
        else:
            auto_cls = AutoModel
        self.vision_model = auto_cls.from_pretrained(
            self.config.vision_name_or_path, config=self.config.vision_config)
        logger.info(
            f"Loaded `{self.vision_model.__class__.__name__}` from {self.config.vision_name_or_path}")

    def init_tokenizer_and_image_processor(self):
        self.tokenizer = AutoTokenizer.from_pretrained(
            self.config.text_name_or_path)
        # SigLIP text model expects sequences of 64 tokens.
        # If using SigLIP, config.tokenizer_max_length must be set to 64.
        # The original SigLIP HF config has model_max_length set to a very high number, so
        # we exploit that to set tokenizer_max_length to 64.
        # In general, we expect 512 tokens at most, as the text model with the longest context that we tried is ColBERTv2 that uses 512.
        if self.tokenizer.model_max_length > 512:
            self.tokenizer.model_max_length = 512 if self.config.tokenizer_max_length is None else self.config.tokenizer_max_length
            logger.info(f"Set tokenizer.model_max_length to {self.tokenizer.model_max_length}")

        self.image_processor = AutoImageProcessor.from_pretrained(
            self.config.vision_name_or_path)

    def freeze(self):
        if self.config.text_frozen:
            self.text_model.requires_grad_(False)
            # CLIP
            if hasattr(self.text_model, 'text_projection'):
                self.text_model.text_projection.requires_grad_(True)
            # SigLIP
            elif hasattr(self.text_model, 'text_model') and hasattr(self.text_model.text_model, 'head'):
                self.text_model.text_model.head.requires_grad_(True)
            # BERT/ColBERT
            elif hasattr(self.text_model, 'pooler'):
                self.text_model.pooler.requires_grad_(True)

        if self.config.vision_frozen:
            self.vision_model.requires_grad_(False)
            # CLIP
            if hasattr(self.vision_model, 'visual_projection'):
                self.vision_model.visual_projection.requires_grad_(True)
            # SigLIP
            elif hasattr(self.vision_model, 'vision_model') and hasattr(self.vision_model.vision_model, 'head'):
                self.vision_model.vision_model.head.requires_grad_(True)

    def __init__(self, config: Ret2Config):
        super().__init__(config)

        self.text_model = None
        self.vision_model = None
        self.init_pretrained_models()

        self.tokenizer = None
        self.image_processor = None
        self.init_tokenizer_and_image_processor()

        self.cell = Ret2Cell(config)

        if config.logit_scale_init_value is None:
            self.logit_scale = None
        else:
            self.logit_scale = nn.Parameter(
                torch.tensor(config.logit_scale_init_value))

        if config.w_final_size:
            self.w_final = nn.Linear(
                config.hidden_size, config.w_final_size, bias=False)
            ret_out_size = config.w_final_size
        else:
            self.w_final = nn.Identity()
            ret_out_size = config.hidden_size

        text_global_size = getattr(config.text_config, 'projection_dim', getattr(
            config.text_config, 'projection_size', config.text_config.hidden_size))
        if text_global_size != ret_out_size:
            self.text_global_feats_proj = nn.Linear(
                text_global_size, ret_out_size, bias=False)
        else:
            self.text_global_feats_proj = nn.Identity()

        vision_global_size = getattr(config.vision_config, 'projection_dim', getattr(
            config.vision_config, 'projection_size', config.vision_config.hidden_size))
        if vision_global_size != ret_out_size:
            self.vision_global_feats_proj = nn.Linear(
                vision_global_size, ret_out_size, bias=False)
        else:
            self.vision_global_feats_proj = nn.Identity()

        self.freeze()

    def train(self, mode: bool = True):
        super().train(mode)
        if self.config.text_frozen:
            self.text_model.eval()
        if self.config.vision_frozen:
            self.vision_model.eval()

    def get_ret_features(
        self,
        pixel_values: Optional[torch.Tensor] = None,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        text_mask: Optional[torch.Tensor] = None,
        vision_mask: Optional[torch.Tensor] = None
    ):
        text_outputs = self.text_model(
            input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)
        vision_outputs = self.vision_model(
            pixel_values=pixel_values, output_hidden_states=True)

        text_hidden_states = torch.stack(
            [text_outputs.hidden_states[i + 1] for i in self.config.text_layer_idxs])
        del text_outputs.hidden_states
        vision_hidden_states = torch.stack(
            [vision_outputs.hidden_states[i + 1] for i in self.config.vision_layer_idxs])
        del vision_outputs.hidden_states

        bsz = input_ids.size(0)
        if attention_mask is not None:
            attention_mask = get_additive_attn_mask(attention_mask, self.dtype, force_attn_on_one_item=text_mask is not None).unsqueeze(
                1).expand(bsz, 1, -1).repeat_interleave(self.cell.text_xa.num_heads, dim=0)
            
        if text_mask is None:
            text_mask = torch.ones(bsz, dtype=self.dtype, device=self.device)
        if vision_mask is None:
            vision_mask = torch.ones(bsz, dtype=self.dtype, device=self.device)

        ret_feats = None
        for (text_hidden_state, vision_hidden_state) in zip(text_hidden_states, vision_hidden_states):
            ret_feats = self.cell(
                hidden_state=ret_feats,
                text_key_value=text_hidden_state,
                text_attention_mask=attention_mask,
                text_mask=text_mask[:, None, None],
                vision_key_value=vision_hidden_state,
                vision_mask=vision_mask[:, None, None]
            )

        ret_feats = self.w_final(ret_feats.squeeze(1))

        text_globfeats = self.text_global_feats_proj(
            text_outputs.get('text_embeds', text_outputs.get('pooler_output')))
        vision_globfeats = self.vision_global_feats_proj(
            vision_outputs.get('image_embeds', vision_outputs.get('pooler_output')))

        return ret_feats +\
            text_globfeats * text_mask[:, None] +\
            vision_globfeats * vision_mask[:, None]

    def forward(
        self,
        query_pixel_values: Optional[torch.Tensor] = None,
        query_text_mask: Optional[torch.Tensor] = None,
        query_vision_mask: Optional[torch.Tensor] = None,
        query_input_ids: Optional[torch.Tensor] = None,
        query_attention_mask: Optional[torch.Tensor] = None,
        passage_pixel_values: Optional[torch.Tensor] = None,
        passage_input_ids: Optional[torch.Tensor] = None,
        passage_attention_mask: Optional[torch.Tensor] = None,
        passage_text_mask: Optional[torch.Tensor] = None,
        passage_vision_mask: Optional[torch.Tensor] = None,
        label_ids: Optional[List[Hashable]] = None
    ):
        query_feats = self.get_ret_features(
            pixel_values=query_pixel_values,
            input_ids=query_input_ids,
            attention_mask=query_attention_mask,
            text_mask=query_text_mask,
            vision_mask=query_vision_mask
        )

        passage_feats = self.get_ret_features(
            pixel_values=passage_pixel_values,
            input_ids=passage_input_ids,
            attention_mask=passage_attention_mask,
            text_mask=passage_text_mask,
            vision_mask=passage_vision_mask
        )

        loss_model = ContrastiveLossModel()

        logit_scale = self.get_logit_scale()
        if logit_scale is not None and self.config.logit_scale_clamp > 0.0:
            logit_scale = logit_scale.clamp(
                max=self.config.logit_scale_clamp)

        loss_outputs = loss_model(
            query_feats=F.normalize(query_feats, p=2, dim=-1),
            passage_feats=F.normalize(passage_feats, p=2, dim=-1),
            logit_scale=logit_scale,
            label_ids=label_ids
        )

        return Ret2ModelOutput(
            query_feats=query_feats.detach(),
            passage_feats=passage_feats.detach(),
            **loss_outputs
        )
