# --------------------------------------------------------
# InternVL
# Copyright (c) 2024 OpenGVLab
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------

import warnings
from typing import List, Optional, Tuple, Union

import torch.utils.checkpoint
import transformers
from torch import nn
from torch.nn import CrossEntropyLoss
from transformers import (AutoModel, GenerationConfig, LlamaForCausalLM,
                          Qwen2ForCausalLM)
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import ModelOutput, logging

from .configuration_internvl_chat import InternVLChatConfig
from .conversation import get_conv_template
from .modeling_intern_vit import InternVisionModel, has_flash_attn
from experiments.models.tactile_projector import TactileProjector
from experiments.models.tactile_encoder import TactileEncoder

logger = logging.get_logger(__name__)


def version_cmp(v1, v2, op='eq'):
    import operator

    from packaging import version
    op_func = getattr(operator, op)
    return op_func(version.parse(v1), version.parse(v2))


class InternVLChatModel(PreTrainedModel):
    config_class = InternVLChatConfig
    main_input_name = 'pixel_values'
    base_model_prefix = 'language_model'
    _supports_flash_attn_2 = True
    supports_gradient_checkpointing = True
    _no_split_modules = ['InternVisionModel', 'LlamaDecoderLayer', 'Qwen2DecoderLayer']

    def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None, use_flash_attn=True):
        super().__init__(config)

        assert version_cmp(transformers.__version__, '4.37.0', 'ge')
        image_size = config.force_image_size or config.vision_config.image_size
        patch_size = config.vision_config.patch_size
        self.patch_size = patch_size
        self.select_layer = config.select_layer
        self.template = config.template
        self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio ** 2))
        self.downsample_ratio = config.downsample_ratio
        self.ps_version = config.ps_version
        use_flash_attn = use_flash_attn if has_flash_attn else False
        config.vision_config.use_flash_attn = True if use_flash_attn else False
        config.llm_config._attn_implementation = 'flash_attention_2' if use_flash_attn else 'eager'

        logger.info(f'num_image_token: {self.num_image_token}')
        logger.info(f'ps_version: {self.ps_version}')
        if vision_model is not None:
            self.vision_model = vision_model
        else:
            self.vision_model = InternVisionModel(config.vision_config)
        if language_model is not None:
            self.language_model = language_model
        else:
            if config.llm_config.architectures[0] == 'LlamaForCausalLM':
                self.language_model = LlamaForCausalLM(config.llm_config)
            elif config.llm_config.architectures[0] == 'Qwen2ForCausalLM':
                self.language_model = Qwen2ForCausalLM(config.llm_config)
            else:
                raise NotImplementedError(f'{config.llm_config.architectures[0]} is not implemented.')
        
        self.tactile_encoder = TactileEncoder()

        vit_hidden_size = config.vision_config.hidden_size
        llm_hidden_size = config.llm_config.hidden_size if hasattr(config, 'llm_config') else config.hidden_size

        self.mlp1 = nn.Sequential(
            nn.LayerNorm(vit_hidden_size * int(1 / self.downsample_ratio) ** 2),
            nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size),
            nn.GELU(),
            nn.Linear(llm_hidden_size, llm_hidden_size)
        )

        self.tactile_projector = TactileProjector(
            in_dim=self.tactile_encoder.embed_dim, # 192
            llm_dim=llm_hidden_size 
        )

        self.img_context_token_id = None
        self.tactile_token_id = 151665
        self.conv_template = get_conv_template(self.template)
        self.system_message = self.conv_template.system_message

    def forward(
            self,
            pixel_values: torch.FloatTensor,
            pixel_values_tactile: torch.FloatTensor | None = None,
            input_ids: torch.LongTensor = None,
            attention_mask: Optional[torch.Tensor] = None,
            position_ids: Optional[torch.LongTensor] = None,
            image_flags: Optional[torch.LongTensor] = None,
            past_key_values: Optional[List[torch.FloatTensor]] = None,
            labels: Optional[torch.LongTensor] = None,
            use_cache: Optional[bool] = None,
            output_attentions: Optional[bool] = None,
            output_hidden_states: Optional[bool] = None,
            return_dict: Optional[bool] = None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        image_flags = image_flags.squeeze(-1)

        # 1️⃣ embedding 가져오기 (clone 제거)
        input_embeds = self.language_model.get_input_embeddings()(input_ids)

        vit_embeds = self.extract_feature(pixel_values)
        vit_embeds = vit_embeds[image_flags == 1]

        B, N, C = input_embeds.shape

        # flatten
        flat_input_ids = input_ids.reshape(-1)
        flat_embeds = input_embeds.reshape(B * N, C)

        # =========================
        # 🔵 Vision injection (연산 기반)
        # =========================
        selected = (flat_input_ids == self.img_context_token_id)

        if selected.sum() > 0:
            vit_flat = vit_embeds.reshape(-1, C)
            flat_embeds[selected] = vit_flat.to(flat_embeds.device)


        # =========================
        # 🟣 Tactile injection (Autograd-safe, shape-safe)
        # =========================
        if pixel_values_tactile is not None:

            tactile_features = self.tactile_encoder(pixel_values_tactile)
            tactile_embeds = self.tactile_projector(tactile_features).to(flat_embeds.dtype)

            tactile_selected = (flat_input_ids == self.tactile_token_id)

            # print(f"🕵️‍♂️ [DEBUG] tactile_token_id: {self.tactile_token_id}")
            # print(f"🕵️‍♂️ [DEBUG] tactile_selected.sum(): {tactile_selected.sum().item()}")

            if tactile_selected.sum() > 0:
                tactile_flat = tactile_embeds.reshape(-1, C)
                selected_indices = tactile_selected.nonzero(as_tuple=False).squeeze(-1)

                # clone 후 scatter (그래프 유지)
                flat_embeds = flat_embeds.clone()

                # flat_embeds = flat_embeds.scatter(
                #     0,
                #     selected_indices.unsqueeze(-1).expand(-1, C),
                #     tactile_flat
                # )
                flat_embeds[selected_indices] = tactile_flat
            #     print(tactile_embeds.requires_grad)
            #     print(flat_embeds.requires_grad)
            else:
                print("🚨 input_ids 안에 촉각 토큰(151665)이 하나도 없습니다!")

            # print("tactile_selected.sum():", tactile_selected.sum())
            # print("tactile_flat.shape[0]:", tactile_flat.shape[0])
        else:
            print("🚨 앗! pixel_values_tactile 데이터가 들어오지 않았습니다 (None)!")
        # reshape back
        input_embeds = flat_embeds.reshape(B, N, C)

        # print("🔍 LLM input requires_grad:", input_embeds.requires_grad)

        outputs = self.language_model(
            inputs_embeds=input_embeds,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        logits = outputs.logits

        loss = None
        if labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # print("shift_logits shape:", shift_logits.shape) 
            # print("shift_labels shape:", shift_labels.shape) 
            # print("shift_logits sample:", shift_logits[0, :10]) 
            # print("shift_labels sample:", shift_labels[0, :10]) 
            # Flatten the tokens
            loss_fct = CrossEntropyLoss()
            shift_logits = shift_logits.view(-1, self.language_model.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:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        if labels is not None:
            flat_labels = labels.reshape(-1)
            tactile_positions = (flat_input_ids == self.tactile_token_id)
            # print("tactile label values:",
            #     flat_labels[tactile_positions][:10])

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

    def pixel_shuffle(self, x, scale_factor=0.5):
        n, w, h, c = x.size()
        # N, W, H, C --> N, W, H * scale, C // scale
        x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))
        # N, W, H * scale, C // scale --> N, H * scale, W, C // scale
        x = x.permute(0, 2, 1, 3).contiguous()
        # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)
        x = x.view(n, int(h * scale_factor), int(w * scale_factor),
                   int(c / (scale_factor * scale_factor)))
        if self.ps_version == 'v1':
            warnings.warn("In ps_version 'v1', the height and width have not been swapped back, "
                          'which results in a transposed image.')
        else:
            x = x.permute(0, 2, 1, 3).contiguous()
        return x

    def extract_feature(self, pixel_values):
        if self.select_layer == -1:
            vit_embeds = self.vision_model(
                pixel_values=pixel_values,
                output_hidden_states=False,
                return_dict=True).last_hidden_state
        else:
            vit_embeds = self.vision_model(
                pixel_values=pixel_values,
                output_hidden_states=True,
                return_dict=True).hidden_states[self.select_layer]
        vit_embeds = vit_embeds[:, 1:, :]

        h = w = int(vit_embeds.shape[1] ** 0.5)
        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
        vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio)
        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
        vit_embeds = self.mlp1(vit_embeds)
        return vit_embeds

    def batch_chat(self, tokenizer, pixel_values, questions, generation_config, num_patches_list=None, pixel_values_tactile=None, num_tactile_tokens_list=None,
                   history=None, return_history=False, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>',
                   TAC_START_TOKEN='<tac>', TAC_END_TOKEN='</tac>', TAC_CONTEXT_TOKEN='<TAC_CONTEXT>', verbose=False, image_counts=None):
        if history is not None or return_history:
            print('Now multi-turn chat is not supported in batch_chat.')
            raise NotImplementedError

        if image_counts is not None:
            num_patches_list = image_counts
            print('Warning: `image_counts` is deprecated. Please use `num_patches_list` instead.')

        img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
        self.img_context_token_id = img_context_token_id

        # Tactile 토큰 ID 자동 매핑 (만약 tokenizer에 없으면 기본값 유지)
        tac_id = tokenizer.convert_tokens_to_ids(TAC_CONTEXT_TOKEN)
        if tac_id != tokenizer.unk_token_id and tac_id is not None:
            self.tactile_token_id = tac_id

        if verbose and pixel_values is not None:
            image_bs = pixel_values.shape[0]
            print(f'dynamic ViT batch size: {image_bs}')

        queries = []
        for idx in range(len(questions)):
            question = questions[idx]
            # 🌟 2. [핵심] 프롬프트 맨 앞에 Vision 토큰 "바로 앞"에 Tactile 토큰 자동 삽입
            prefix = ""
            if pixel_values_tactile is not None and '<tactile>' not in question:
                prefix += '<tactile>\n'
            if pixel_values is not None and '<image>' not in question:
                prefix += '<image>\n'
            question = prefix + question # 결과: "<tactile>\n<image>\n질문내용"

            template = get_conv_template(self.template)
            template.system_message = self.system_message
            template.append_message(template.roles[0], question)
            template.append_message(template.roles[1], None)
            query = template.get_prompt()

            # 🌟 3. <tactile> 문자열을 실제 Context 토큰 길이만큼 확장
            # 🌟 3. <tactile> 확장
            if num_tactile_tokens_list is not None and len(num_tactile_tokens_list) > idx:
                num_tac = num_tactile_tokens_list[idx]
                tac_tokens = TAC_START_TOKEN + TAC_CONTEXT_TOKEN * num_tac + TAC_END_TOKEN
                query = query.replace('<tactile>', tac_tokens, 1)

            # 4. <image> 확장 (num_patches_list가 있을 때만 작동하도록 안전장치)
            if num_patches_list is not None and len(num_patches_list) > idx:
                num_patches = num_patches_list[idx]
                image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
                query = query.replace('<image>', image_tokens, 1)

            queries.append(query)

        tokenizer.padding_side = 'left'
        model_inputs = tokenizer(queries, return_tensors='pt', padding=True)
        input_ids = model_inputs['input_ids'].to(self.device)
        attention_mask = model_inputs['attention_mask'].to(self.device)
        eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())
        generation_config['eos_token_id'] = eos_token_id
        generation_output = self.generate(
            pixel_values=pixel_values,
            pixel_values_tactile=pixel_values_tactile,
            input_ids=input_ids,
            attention_mask=attention_mask,
            **generation_config
        )
        responses = tokenizer.batch_decode(generation_output, skip_special_tokens=True)
        responses = [response.split(template.sep.strip())[0].strip() for response in responses]
        return responses

    def chat(self, tokenizer, pixel_values, question, generation_config, pixel_values_tactile=None, num_tactile_tokens_list=None, history=None, return_history=False,
             num_patches_list=None, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>', TAC_START_TOKEN='<tac>', TAC_END_TOKEN='</tac>', TAC_CONTEXT_TOKEN='<TAC_CONTEXT>',
             verbose=False):

        # 🌟 1. [핵심] Vision 토큰 바로 앞에 Tactile 토큰 자동 삽입
        if history is None:
            prefix = ""
            if pixel_values_tactile is not None and '<tactile>' not in question:
                prefix += '<tactile>\n'
            if pixel_values is not None and '<image>' not in question:
                prefix += '<image>\n'
            question = prefix + question

        if num_patches_list is None:
            num_patches_list = [pixel_values.shape[0]] if pixel_values is not None else []
        assert pixel_values is None or len(pixel_values) == sum(num_patches_list)

        # 2. 토큰 ID 등록
        self.img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
        tac_id = tokenizer.convert_tokens_to_ids(TAC_CONTEXT_TOKEN)
        if tac_id != tokenizer.unk_token_id and tac_id is not None:
            self.tactile_token_id = tac_id

        template = get_conv_template(self.template)
        template.system_message = self.system_message
        eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())

        history = [] if history is None else history
        for (old_question, old_answer) in history:
            template.append_message(template.roles[0], old_question)
            template.append_message(template.roles[1], old_answer)
        template.append_message(template.roles[0], question)
        template.append_message(template.roles[1], None)
        query = template.get_prompt()

        if verbose and pixel_values is not None:
            print(f'dynamic ViT batch size: {pixel_values.shape[0]}')

        # 🌟 3. <tactile> 토큰 확장
        if num_tactile_tokens_list is not None:
            for num_tac in num_tactile_tokens_list:
                tac_tokens = TAC_START_TOKEN + TAC_CONTEXT_TOKEN * num_tac + TAC_END_TOKEN
                query = query.replace('<tactile>', tac_tokens, 1)

        # 4. <image> 토큰 확장
        for num_patches in num_patches_list:
            image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
            query = query.replace('<image>', image_tokens, 1)

        model_inputs = tokenizer(query, return_tensors='pt')
        input_ids = model_inputs['input_ids'].to(self.device)
        attention_mask = model_inputs['attention_mask'].to(self.device)
        generation_config['eos_token_id'] = eos_token_id
        generation_output = self.generate(
            pixel_values=pixel_values,
            pixel_values_tactile=pixel_values_tactile,
            input_ids=input_ids,
            attention_mask=attention_mask,
            **generation_config
        )
        response = tokenizer.batch_decode(generation_output, skip_special_tokens=True)[0]
        response = response.split(template.sep.strip())[0].strip()
        history.append((question, response))
        if return_history:
            return response, history
        else:
            query_to_print = query.replace(IMG_CONTEXT_TOKEN, '')
            query_to_print = query_to_print.replace(TAC_CONTEXT_TOKEN, '')
            query_to_print = query_to_print.replace(f'{IMG_START_TOKEN}{IMG_END_TOKEN}', '<image>')
            query_to_print = query_to_print.replace(f'{TAC_START_TOKEN}{TAC_END_TOKEN}', '<tactile>')
            if verbose:
                print(query_to_print, response)
            return response
    
    @property
    def all_tied_weights_keys(self):
        keys = getattr(self, "_tied_weights_keys", None)
        if keys is None:
            return {}
        return keys

    @torch.no_grad()
    def generate(
            self,
            pixel_values: Optional[torch.FloatTensor] = None,
            pixel_values_tactile: Optional[torch.FloatTensor] = None,
            input_ids: Optional[torch.FloatTensor] = None,
            attention_mask: Optional[torch.LongTensor] = None,
            visual_features: Optional[torch.FloatTensor] = None,
            generation_config: Optional[GenerationConfig] = None,
            output_hidden_states: Optional[bool] = None,
            **generate_kwargs,
    ) -> torch.LongTensor:

        assert self.img_context_token_id is not None
        if pixel_values is not None:
            if visual_features is not None:
                vit_embeds = visual_features
            else:
                vit_embeds = self.extract_feature(pixel_values)
            input_embeds = self.language_model.get_input_embeddings()(input_ids)
            B, N, C = input_embeds.shape
            input_embeds = input_embeds.reshape(B * N, C)

            input_ids = input_ids.reshape(B * N)
            selected = (input_ids == self.img_context_token_id)
            assert selected.sum() != 0
            input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device)

            input_embeds = input_embeds.reshape(B, N, C)
        else:
            input_embeds = self.language_model.get_input_embeddings()(input_ids)

        if pixel_values_tactile is not None:
            # 인코더 및 프로젝터 통과
            tactile_features = self.tactile_encoder(pixel_values_tactile)
            tactile_features.requires_grad_(True)
            print("tactile_features.requires_grad:", tactile_features.requires_grad)
            tactile_embeds = self.tactile_projector(tactile_features).to(input_embeds.dtype)

            # print("tactile_embeds.requires_grad:", tactile_embeds.requires_grad)
            
            # 위쪽 비전 코드 통과 여부에 상관없이 다시 한번 3D 형태로 펴서 작업
            B, N, C = input_embeds.shape
            input_embeds = input_embeds.reshape(B * N, C)
            
            # input_ids가 1차원(B*N)이든 2차원(B,N)이든 관계없이 1차원으로 펴서 비교
            flat_input_ids = input_ids.reshape(-1)
            
            tactile_selected = (flat_input_ids == self.tactile_token_id)
            
            # 만약 입력 프롬프트에 촉각 토큰이 들어있다면 덮어씌움
            if tactile_selected.sum() > 0:
                input_embeds[tactile_selected] = tactile_embeds.reshape(-1, C).to(input_embeds.device)
            
            # 다시 LLM에 들어갈 3D 형태로 원상 복구
            input_embeds = input_embeds.reshape(B, N, C)
        # ==========================================================

        # print("🔍 LLM input requires_grad:", input_embeds.requires_grad)

        # print("tactile_selected.sum():", tactile_selected.sum().item())
        # print("tactile_embed_tokens:", tactile_embeds.reshape(-1, C).shape[0])
        # print("vision norm:", vit_embeds.norm().item())
        # print("tactile norm:", tactile_embeds.norm().item())
        print("vision token avg norm:", vit_embeds.norm(dim=-1).mean().item())
    # print("tactile token avg norm:", tactile_embeds.norm(dim=-1).mean().item())


        outputs = self.language_model.generate(
            inputs_embeds=input_embeds,
            attention_mask=attention_mask,
            generation_config=generation_config,
            output_hidden_states=output_hidden_states,
            use_cache=True,
            **generate_kwargs,
        )

        return outputs

    @property
    def lm_head(self):
        return self.language_model.get_output_embeddings()

    def get_input_embeddings(self):
        return self.language_model.get_input_embeddings()

    def get_output_embeddings(self):
        return self.language_model.get_output_embeddings()