import gc
import math
from typing import List, Optional, Tuple, Union

import torch
import torch.nn as nn
from torch import Tensor
from torch.nn import CrossEntropyLoss
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    MistralConfig,
    MistralForCausalLM,
    MistralModel,
)
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
)

from ..model.resampler import Resampler
from ..model.utils import build_transform

DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"
DEFAULT_IM_START_TOKEN = "<im_start>"
DEFAULT_IM_END_TOKEN = "<im_end>"


class OmniLMMConfig(MistralConfig):
    model_type = "omnilmm"


class Identity(torch.nn.Identity):
    def forward(self, input: Tensor, **kwargs) -> Tensor:
        return super().forward(input)


def create_vision_module(config):
    import timm

    vision_tower = timm.create_model(
        "eva02_enormous_patch14_clip_224.laion2b_plus",
        pretrained=False,
        num_classes=0,
        dynamic_img_size=True,
        dynamic_img_pad=True,
    )

    if isinstance(vision_tower, timm.models.VisionTransformer):
        if vision_tower.attn_pool is not None:
            vision_tower.attn_pool = Identity()

    # use 2nd last layer's output
    vision_tower.blocks[-1] = Identity()

    embed_dim = config.hidden_size
    resampler = Resampler(
        grid_size=int(math.sqrt(config.num_query)),
        embed_dim=embed_dim,
        num_heads=embed_dim // 128,
        kv_dim=vision_tower.embed_dim,
    )
    return vision_tower, resampler


class OmniLMMModel(MistralModel):
    config_class = OmniLMMConfig

    def __init__(
        self,
        config: OmniLMMConfig,
        mm_vision_tower=None,
        mm_hidden_size=None,
        tune_clip=True,
    ):
        super(OmniLMMModel, self).__init__(config)

        if hasattr(config, "mm_vision_tower"):
            vision_tower, resampler = create_vision_module(config)

            # print(__file__, 'skip loading vision tower weights')

            # HACK: for FSDP
            self.vision_tower = [vision_tower]
            self.resampler = resampler
            if tune_clip:
                self.vision_tower = self.vision_tower[0]

        self.vision_config = lambda x: None

    def initialize_vision_modules(
        self, vision_tower, no_randaug, num_query, image_size, tune_clip=False
    ):
        self.config.mm_vision_tower = vision_tower
        self.config.use_mm_proj = True
        self.config.num_query = num_query
        self.config.image_size = image_size

        if not hasattr(self, "vision_tower"):
            vision_tower, resampler = create_vision_module(self.config)
            state_dict = torch.load(
                "/tt/data/public/multimodal/multimodal_model_ckpts/timm/eva02_enormous_patch14_clip_224.laion2b_plus.pt"
            )
            vision_tower.load_state_dict(state_dict, strict=False)
            del state_dict
            gc.collect()
        else:
            if isinstance(self.vision_tower, list):
                vision_tower = self.vision_tower[0]
            else:
                vision_tower = self.vision_tower
            resampler = self.resampler
        self.vision_tower = vision_tower if tune_clip else [vision_tower]
        self.resampler = resampler

        train_img_transform = build_transform(
            is_train=True,
            randaug=not no_randaug,
            input_size=self.config.image_size,
            std_mode="OPENAI_CLIP",
        )
        eval_img_transform = build_transform(
            is_train=False, input_size=self.config.image_size, std_mode="OPENAI_CLIP"
        )

        return dict(
            image_processor=(train_img_transform, eval_img_transform),
            image_token_len=num_query,
            vision_config=self.vision_config,
        )

    def get_vision_embedding(self, pixel_values):
        if isinstance(self.vision_tower, list):
            vision_tower = self.vision_tower[0]  # HACK: for FSDP
        else:
            vision_tower = self.vision_tower

        dtype = vision_tower.pos_embed.data.dtype
        vision_embedding = vision_tower.forward_features(pixel_values.type(dtype))
        if (
            hasattr(vision_tower, "num_prefix_tokens")
            and vision_tower.num_prefix_tokens > 0
        ):
            vision_embedding = vision_embedding[:, vision_tower.num_prefix_tokens :]
        res = self.resampler(vision_embedding)
        return res

    def get_vllm_embedding(self, data):
        if "vision_hidden_states" not in data:
            pixel_values_list = data["pixel_values"]
            vision_hidden_states = []
            for pixel_values in pixel_values_list:
                if len(pixel_values) > 0:
                    vision_hidden_states.append(
                        self.get_vision_embedding(pixel_values.unsqueeze(0))[0]
                    )
                else:
                    vision_hidden_states.append([])
        else:
            vision_hidden_states = data["vision_hidden_states"]

        # vllm_embedding = self.llm.model.embed_tokens(data['input_ids']) * self.llm.config.scale_emb
        inputs_embeds = self.embed_tokens(data["input_ids"])
        vision_hidden_states = [
            i.type(inputs_embeds.dtype) if isinstance(i, torch.Tensor) else i
            for i in vision_hidden_states
        ]

        # HACK: replace back original embeddings for LLaVA pretraining
        orig_embeds_params = getattr(self, "orig_embeds_params", None)

        new_input_embeds = []
        cur_image_idx = 0
        for cur_input_ids, cur_input_embeds in zip(data["input_ids"], inputs_embeds):
            if (cur_input_ids == self.vision_config.im_patch_token).sum() == 0:
                # multimodal LLM, but the current sample is not multimodal
                cur_input_embeds = cur_input_embeds + (0.0 * dummy_image_features).sum()
                new_input_embeds.append(cur_input_embeds)
                continue

            if self.vision_config.use_im_start_end:
                cur_image_features = vision_hidden_states[cur_image_idx]
                num_patches = cur_image_features.shape[0]
                if (cur_input_ids == self.vision_config.im_start_token).sum() != (
                    cur_input_ids == self.vision_config.im_end_token
                ).sum():
                    raise ValueError(
                        "The number of image start tokens and image end tokens should be the same."
                    )
                image_start_tokens = torch.where(
                    cur_input_ids == self.vision_config.im_start_token
                )[0]
                for image_start_token_pos in image_start_tokens:
                    cur_image_features = vision_hidden_states[cur_image_idx].to(
                        device=cur_input_embeds.device
                    )
                    num_patches = cur_image_features.shape[0]
                    if (
                        cur_input_ids[image_start_token_pos + num_patches + 1]
                        != self.vision_config.im_end_token
                    ):
                        raise ValueError(
                            "The image end token should follow the image start token."
                        )
                    if orig_embeds_params is not None:
                        cur_new_input_embeds = torch.cat(
                            (
                                cur_input_embeds[:image_start_token_pos].detach(),
                                cur_input_embeds[
                                    image_start_token_pos : image_start_token_pos + 1
                                ],
                                cur_image_features,
                                cur_input_embeds[
                                    image_start_token_pos
                                    + num_patches
                                    + 1 : image_start_token_pos
                                    + num_patches
                                    + 2
                                ],
                                cur_input_embeds[
                                    image_start_token_pos + num_patches + 2 :
                                ].detach(),
                            ),
                            dim=0,
                        )
                    else:
                        cur_new_input_embeds = torch.cat(
                            (
                                cur_input_embeds[: image_start_token_pos + 1],
                                cur_image_features,
                                cur_input_embeds[
                                    image_start_token_pos + num_patches + 1 :
                                ],
                            ),
                            dim=0,
                        )
                    cur_image_idx += 1
                new_input_embeds.append(cur_new_input_embeds)
            else:
                raise NotImplementedError
        inputs_embeds = torch.stack(new_input_embeds, dim=0)

        return inputs_embeds, vision_hidden_states

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = 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,
        **kwargs,
    ) -> Union[Tuple, BaseModelOutputWithPast]:
        # HACK: replace back original embeddings for LLaVA pretraining
        orig_embeds_params = getattr(self, "orig_embeds_params", None)

        if inputs_embeds is None and past_key_values is None:
            inputs_embeds = self.embed_tokens(input_ids)

            vision_tower = getattr(self, "vision_tower", None)
            if (
                vision_tower is not None
                and (input_ids.shape[1] != 1 or self.training)
                and images is not None
            ):
                if type(images) is list:
                    image_features = []
                    for image in images:
                        image_forward_out = self.get_vision_embedding(
                            image.unsqueeze(0)
                        )[0]
                        image_features.append(image_forward_out)
                else:
                    image_features = self.get_vision_embedding(images)

                dummy_image_features = torch.zeros(
                    self.config.num_query,
                    self.config.hidden_size,
                    device=inputs_embeds.device,
                    dtype=inputs_embeds.dtype,
                )

                new_input_embeds = []
                cur_image_idx = 0
                for cur_input_ids, cur_input_embeds in zip(input_ids, inputs_embeds):
                    if (cur_input_ids == self.vision_config.im_patch_token).sum() == 0:
                        # multimodal LLM, but the current sample is not multimodal
                        cur_input_embeds = (
                            cur_input_embeds + (0.0 * dummy_image_features).sum()
                        )
                        new_input_embeds.append(cur_input_embeds)
                        continue

                    if self.vision_config.use_im_start_end:
                        cur_image_features = image_features[cur_image_idx]
                        num_patches = cur_image_features.shape[0]
                        if (
                            cur_input_ids == self.vision_config.im_start_token
                        ).sum() != (
                            cur_input_ids == self.vision_config.im_end_token
                        ).sum():
                            raise ValueError(
                                "The number of image start tokens and image end tokens should be the same."
                            )
                        image_start_tokens = torch.where(
                            cur_input_ids == self.vision_config.im_start_token
                        )[0]
                        for image_start_token_pos in image_start_tokens:
                            cur_image_features = image_features[cur_image_idx].to(
                                device=cur_input_embeds.device
                            )
                            num_patches = cur_image_features.shape[0]
                            if (
                                cur_input_ids[image_start_token_pos + num_patches + 1]
                                != self.vision_config.im_end_token
                            ):
                                raise ValueError(
                                    "The image end token should follow the image start token."
                                )
                            if orig_embeds_params is not None:
                                cur_new_input_embeds = torch.cat(
                                    (
                                        cur_input_embeds[
                                            :image_start_token_pos
                                        ].detach(),
                                        cur_input_embeds[
                                            image_start_token_pos : image_start_token_pos
                                            + 1
                                        ],
                                        cur_image_features,
                                        cur_input_embeds[
                                            image_start_token_pos
                                            + num_patches
                                            + 1 : image_start_token_pos
                                            + num_patches
                                            + 2
                                        ],
                                        cur_input_embeds[
                                            image_start_token_pos + num_patches + 2 :
                                        ].detach(),
                                    ),
                                    dim=0,
                                )
                            else:
                                cur_new_input_embeds = torch.cat(
                                    (
                                        cur_input_embeds[: image_start_token_pos + 1],
                                        cur_image_features,
                                        cur_input_embeds[
                                            image_start_token_pos + num_patches + 1 :
                                        ],
                                    ),
                                    dim=0,
                                )
                            cur_image_idx += 1
                        new_input_embeds.append(cur_new_input_embeds)
                    else:
                        raise NotImplementedError
                inputs_embeds = torch.stack(new_input_embeds, dim=0)
                input_ids = None

        return super(OmniLMMModel, self).forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            **kwargs,
        )


class OmniLMMForCausalLM(MistralForCausalLM):
    config_class = OmniLMMConfig

    def __init__(self, config, mm_vision_tower=None, tune_clip=True):
        super(MistralForCausalLM, self).__init__(config)
        self.model = OmniLMMModel(
            config, mm_vision_tower=mm_vision_tower, tune_clip=tune_clip
        )

        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = 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,
        **kwargs,
    ) -> 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
        )

        # print(f'@@@ At forward, labels: {labels.shape}-{labels}', flush=True)
        # print(f'@@@ At forward, input_ids: {input_ids.shape}-{input_ids}', flush=True)
        # print(f'@@@ At forward, input_ids: {attention_mask.shape}-{attention_mask}', flush=True)

        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            images=images,
            **kwargs,
        )

        hidden_states = outputs[0]
        logits = self.lm_head(hidden_states)

        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/pipeline parallelism
            shift_labels = shift_labels.to(shift_logits.device)
            loss = loss_fct(shift_logits, shift_labels)

        if not return_dict:
            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,
        )

    # TODO could be removed for generate_vllm()
    def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        **kwargs,
    ):
        if past_key_values:
            input_ids = input_ids[:, -1:]

        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
        if inputs_embeds is not None and past_key_values is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}

        model_inputs.update(
            {
                "past_key_values": past_key_values,
                "use_cache": kwargs.get("use_cache"),
                "attention_mask": attention_mask,
                "images": kwargs.get("images", None),
            }
        )
        return model_inputs

    def generate_vllm(
        self,
        input_ids: torch.LongTensor = None,
        images: Optional[torch.FloatTensor] = None,
        vision_hidden_states=None,
        return_vision_hidden_states=False,
        **kwargs,
    ):
        model_inputs = {"input_ids": input_ids}
        if vision_hidden_states is None:
            model_inputs["pixel_values"] = images
        else:
            model_inputs["vision_hidden_states"] = vision_hidden_states

        with torch.inference_mode():
            inputs_embeds, vision_hidden_states = self.model.get_vllm_embedding(
                model_inputs
            )

            result = self.generate(inputs_embeds=inputs_embeds, **kwargs)

        if return_vision_hidden_states:
            return result, vision_hidden_states

        return result

    def initialize_vision_tokenizer(
        self, mm_use_im_start_end, tokenizer, device, tune_mm_mlp_adapter=False
    ):
        self.model.vision_config.use_im_start_end = mm_use_im_start_end
        tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
        self.resize_token_embeddings(len(tokenizer))

        if mm_use_im_start_end:
            num_new_tokens = tokenizer.add_tokens(
                [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True
            )
            self.resize_token_embeddings(len(tokenizer))
            (
                self.model.vision_config.im_start_token,
                self.model.vision_config.im_end_token,
            ) = tokenizer.convert_tokens_to_ids(
                [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN]
            )

            if num_new_tokens > 0:
                input_embeddings = self.get_input_embeddings().weight.data
                output_embeddings = self.get_output_embeddings().weight.data

                input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
                    dim=0, keepdim=True
                )
                output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
                    dim=0, keepdim=True
                )

                input_embeddings[-num_new_tokens:] = input_embeddings_avg
                output_embeddings[-num_new_tokens:] = output_embeddings_avg

            # for new sft data
            num_new_tokens = tokenizer.add_tokens(
                ["<box>", "</box>", "<ref>", "</ref>", "<quad>", "</quad>"],
                special_tokens=True,
            )
            self.resize_token_embeddings(len(tokenizer))

            if num_new_tokens > 0:
                input_embeddings = self.get_input_embeddings().weight.data
                output_embeddings = self.get_output_embeddings().weight.data

                input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
                    dim=0, keepdim=True
                )
                output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
                    dim=0, keepdim=True
                )

                input_embeddings[-num_new_tokens:] = input_embeddings_avg
                output_embeddings[-num_new_tokens:] = output_embeddings_avg

            if tune_mm_mlp_adapter:
                self.model.orig_embeds_params = [
                    self.get_input_embeddings().weight.data.clone().to(device=device)
                ]
                for p in self.get_input_embeddings().parameters():
                    p.requires_grad = True
                for p in self.get_output_embeddings().parameters():
                    p.requires_grad = False

        self.model.vision_config.im_patch_token = tokenizer.convert_tokens_to_ids(
            [DEFAULT_IMAGE_PATCH_TOKEN]
        )[0]
        print(
            f"Tokenizer: {tokenizer}\n patch_token_id: {self.model.vision_config.im_patch_token}, visoin_config: {self.model.vision_config}",
            flush=True,
        )
        # exit()


AutoConfig.register("omnilmm", OmniLMMConfig)
AutoModelForCausalLM.register(OmniLMMConfig, OmniLMMForCausalLM)
