# Copyright 2024 Zhenwei Shao and MILVLG team.
# Licensed under the Apache License, Version 2.0.


from typing import List, Optional, Tuple, Union

import torch
import torch.nn as nn

from transformers import AutoConfig, AutoModelForCausalLM

from .phi2.modeling_phi import PhiConfig, PhiModel, PhiForCausalLM,PhiPreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
from ..llava_arch import LlavaMetaModel, LlavaMetaForCausalLM

class FlashSlothConfig(PhiConfig):
    model_type = "flashsloth"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.image_token_index = getattr(self, "image_token_index", 50297)
        self.image_token = getattr(self, "image_token", "<image>")


class FlashSlothModel(LlavaMetaModel, PhiModel):
    config_class = FlashSlothConfig

    def __init__(self, config: FlashSlothConfig):
        super(FlashSlothModel, self).__init__(config)


class FlashSlothForCausalLM(PhiPreTrainedModel, LlavaMetaForCausalLM):
    """FlashSloth for Causal Language Modeling."""

    # _keys_to_ignore_on_load_missing = [""]
    # _keys_to_ignore_on_load_unexpected = [r"transformer\.h\.\d+\.mlp.(fc_in|fc_out)\.(weight|bias)"]
    config_class = FlashSlothConfig

    def __init__(self, config: FlashSlothConfig) -> None:
        super().__init__(config)

        self.model = FlashSlothModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=True)
        config =self.config
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self) -> nn.Linear:
        return self.lm_head

    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
        self.lm_head = new_embeddings

    def get_model(self):
        return self.model

    def get_decoder(self):
        return self.model
    
    def set_decoder(self, decoder):
        self.model = decoder
    
    def image_preprocess(self, images):
        return self.get_vision_tower().image_processor(images)['pixel_values']

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        images: Optional[torch.FloatTensor] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        learnable_tokens = self.model.get_learnabletoken()
        if inputs_embeds is None:
            (
                input_ids,
                position_ids,
                attention_mask,
                past_key_values,
                inputs_embeds,
                labels,
                insert_place, 
                image_features, 
                learnable_token_len,
                modal,
                question_token_ranges

            ) = self.prepare_inputs_labels_for_multimodal(
                input_ids,
                position_ids,
                attention_mask,
                past_key_values,
                labels,
                images,
                learnable_tokens,
                'phi2',

            )

        outputs = self.model(
            input_ids=input_ids,
            past_key_values=past_key_values, 
            attention_mask=attention_mask,
            position_ids=position_ids, 
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            insert_place=insert_place,
            image_features=image_features,
            learnable_token_len=learnable_token_len,
            modal = modal,
            question_token_ranges = question_token_ranges
            )
        hidden_states = outputs[0]
        logits = self.lm_head(hidden_states)
        logits = logits.float()

        loss = None
        if labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = CrossEntropyLoss()
            shift_logits = shift_logits.view(-1, self.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            # Enable model parallelism
            shift_labels = shift_labels.to(shift_logits.device)
            loss = loss_fct(shift_logits, shift_labels)
        if not return_dict:
            loss = None
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
        images = kwargs.pop("images", None)
        _inputs = super().prepare_inputs_for_generation(
            input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, **kwargs
        )
        if images is not None:
            _inputs['images'] = images
        return _inputs

AutoConfig.register("flashsloth", FlashSlothConfig)
AutoModelForCausalLM.register(FlashSlothConfig, FlashSlothForCausalLM)