# --------------------------------------------------------
# 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 PIL import Image, ImageDraw, ImageFont
import numpy as np
import cv2
import imageio
from scipy.ndimage import gaussian_filter
from PIL import Image, ImageDraw, ImageFont
import tqdm
import random

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))

def draw_text_to_image(text, font, image_width=500, min_height=500, bg_color=(255, 255, 255)):
    paragraphs = text.split('\n')
    # Danh sách chứa tất cả các dòng văn bản sau khi được xử lý
    lines = []
    total_height = 0
    for paragraph in paragraphs:
        words = paragraph.split(' ')
        current_line = ""
        for word in words:
            test_line = current_line + word + " "
            bbox = font.getbbox(test_line)
            width = bbox[2] - bbox[0]
            if width <= image_width - 20:  # Trừ lề khoảng 10px mỗi bên
                current_line = test_line
            else:
                lines.append(current_line)
                current_line = word + " "
                total_height += font.getbbox(current_line)[3]
        lines.append(current_line)  # Thêm dòng cuối cùng của đoạn văn
        total_height += font.getbbox(current_line)[3]
    total_height = int(total_height*1.25)
    if total_height < min_height:
        total_height = min_height
    image = Image.new('RGB', (image_width, total_height), color=bg_color)
    draw = ImageDraw.Draw(image)
    # Vẽ đoạn văn bản tiếng Việt lên ảnh, từng dòng một
    text_color = tuple(random.randint(0, 1) for _ in range(3))
    y_text = 10
    for line in lines:
        draw.text((10, y_text), line, font=font, fill=text_color)
        y_text += font.getbbox(line)[3] * 1.2
    return image

def load_image_v2(image_file, input_size=448, max_num=12):
    image = Image.open(image_file).convert('RGB')
    transform = build_transform(input_size=input_size)
    images, target_aspect_ratio = dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)
    pixel_values = [transform(image) for image in images]
    pixel_values = torch.stack(pixel_values)
    return pixel_values, target_aspect_ratio

def adjust_overlay(overlay, text_img):
    h_o, w_o = overlay.shape[:2]
    h_t, w_t = text_img.shape[:2]

    if h_o > w_o:  # Overlay là ảnh đứng
        # Resize overlay sao cho h = h_t, giữ nguyên tỷ lệ
        new_h = h_t
        new_w = int(w_o * (new_h / h_o))
        overlay_resized = cv2.resize(overlay, (new_w, new_h))
    else:  # Overlay là ảnh ngang
        # Giữ nguyên overlay, nhưng nếu h < h_t thì thêm padding trắng
        overlay_resized = overlay.copy()

    # Thêm padding trắng nếu overlay có h < h_t
    if overlay_resized.shape[0] < h_t:
        pad_h = h_t - overlay_resized.shape[0]
        padding = np.ones((pad_h, overlay_resized.shape[1], 3), dtype=np.uint8) * 255
        overlay_resized = np.vstack((overlay_resized, padding))  # Padding vào dưới

    # Đảm bảo overlay có cùng chiều cao với text_img
    if overlay_resized.shape[0] != h_t:
        overlay_resized = cv2.resize(overlay_resized, (overlay_resized.shape[1], h_t))

    return overlay_resized
        
class InternVLChatModel(PreTrainedModel):
    config_class = InternVLChatConfig
    main_input_name = 'pixel_values'
    base_model_prefix = 'language_model'
    _supports_flash_attn_2 = 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.')

        vit_hidden_size = config.vision_config.hidden_size
        llm_hidden_size = config.llm_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.img_context_token_id = None
        self.conv_template = get_conv_template(self.template)
        self.system_message = self.conv_template.system_message

    def forward(
            self,
            pixel_values: torch.FloatTensor,
            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)
        input_embeds = self.language_model.get_input_embeddings()(input_ids).clone()

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

        B, N, C = input_embeds.shape
        input_embeds = input_embeds.reshape(B * N, C)

        if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:
            print(f'dynamic ViT batch size: {vit_batch_size}, images per sample: {vit_batch_size / B}, dynamic token length: {N}')

        input_ids = input_ids.reshape(B * N)
        selected = (input_ids == self.img_context_token_id)
        try:
            input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)
        except Exception as e:
            vit_embeds = vit_embeds.reshape(-1, C)
            print(f'warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, '
                  f'vit_embeds.shape={vit_embeds.shape}')
            n_token = selected.sum()
            input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds[:n_token]

        input_embeds = input_embeds.reshape(B, N, C)

        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()
            # 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

        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,
                   history=None, return_history=False, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>',
                   IMG_CONTEXT_TOKEN='<IMG_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

        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, num_patches in enumerate(num_patches_list):
            question = questions[idx]
            if pixel_values is not None and '<image>' not in question:
                question = '<image>\n' + question
            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()

            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,
            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, history=None, return_history=False,
             num_patches_list=None, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>',
             verbose=False, attention_visualize=False,last_visualize_layers=7,raw_image_path="",target_aspect_ratio=(1,1)):

        if history is None and pixel_values is not None and '<image>' not in question:
            question = '<image>\n' + 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)

        img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
        self.img_context_token_id = img_context_token_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:
            image_bs = pixel_values.shape[0]
            print(f'dynamic ViT batch size: {image_bs}')

        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
        if attention_visualize:
            generation_output = self.generate(
                pixel_values=pixel_values,
                input_ids=input_ids,
                attention_mask=attention_mask,
                attention_visualize=attention_visualize,
                output_hidden_states=True,
                **generation_config
            )
            return generation_output, query
            #################################### Attention visualize ##################################################
            # attentions_tensors = []
            # for tok_ in generation_output["attentions"]:
            #     attentions_tensors.append([])
            #     for lay_ in tok_ :
            #         attentions_tensors[-1].append(lay_.detach().cpu().type(torch.float).numpy())
            # attention_scores = attentions_tensors
            # query_ = tokenizer(query)
            # start_img_token_index = int(np.where(np.array(query_["input_ids"])==tokenizer("<img>")["input_ids"][0])[0]+1)
            # end_img_token_index = int(np.where(np.array(query_["input_ids"])==tokenizer("</img>")["input_ids"][0])[0]-256)
            # if end_img_token_index - start_img_token_index  == 0 :
            #     end_img_token_index = int(np.where(np.array(query_["input_ids"])==tokenizer("</img>")["input_ids"][0])[0])
            
            # # Đọc ảnh gốc
            # image = cv2.imread(raw_image_path)
            # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            # # Resize ảnh nhỏ hơn để giảm dung lượng GIF
            # scale_factor = 1.  # Giảm 50% kích thước
            # # Font chữ
            # font = ImageFont.truetype("DejaVuSans.ttf", 15)
            # alpha = 0.4  
            # # Lưu danh sách frames GIF
            # visualization_frames = []
            # # Chuỗi sinh ra
            # generated_text = ""
            # frame_step = 1
            # # Lặp qua từng token
            # for index_focus in tqdm.tqdm(range(0, generation_output.sequences.shape[1], frame_step)):
            #     token_text = tokenizer.decode(generation_output.sequences[0, index_focus])
            #     generated_text += token_text  # Ghép chữ lại
            #     # Tạo heatmap trung bình từ các lớp attention
            #     heat_maps = []
            #     for i in range(1, 8):
            #         heat_maps.append(
            #             self.visualize_attention(
            #                 attention_scores[index_focus], layer=-i, head=None,
            #                 start_img_token_index=start_img_token_index, end_img_token_index=end_img_token_index, target_aspect_ratio=target_aspect_ratio
            #             )[0]
            #         )
            #     heatmap = np.array(heat_maps).mean(0)
            #     # Resize heatmap về kích thước ảnh gốc
            #     heatmap = cv2.resize(heatmap, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
            #     # Làm mượt heatmap
            #     heatmap_smooth = gaussian_filter(heatmap, sigma=1)
            #     # Chuẩn hóa heatmap về 0-255
            #     heatmap_norm = cv2.normalize(heatmap_smooth, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
            #     heatmap_color = cv2.applyColorMap(heatmap_norm, cv2.COLORMAP_JET)
            #     heatmap_color = cv2.cvtColor(heatmap_color, cv2.COLOR_BGR2RGB)
            #     # Overlay ảnh heatmap lên ảnh gốc
            #     overlay = cv2.addWeighted(image, 1 - alpha, heatmap_color, alpha, 0)
            #     # Tạo ảnh chứa text bên phải
            #     text_img = draw_text_to_image(generated_text, font, image_width=600, min_height=500)
            #     text_img = np.array(text_img)
            #     # text_img =  cv2.resize(np.array(text_img),(overlay.shape[1],overlay.shape[0]))
            #     # combined_image = np.hstack((overlay, text_img))
            #     ## Đảm bảo overlay và text_img có cùng kích thước
            #     overlay_adjusted = adjust_overlay(overlay, text_img)
            #     # Ghép ảnh
            #     combined_image = np.hstack((overlay_adjusted, text_img))
            #     # Lưu vào danh sách frames
            #     visualization_frames.append(combined_image)

            # generation_output = generation_output.sequences
            # 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, visualization_frames
            # else:
            #     query_to_print = query.replace(IMG_CONTEXT_TOKEN, '')
            #     query_to_print = query_to_print.replace(f'{IMG_START_TOKEN}{IMG_END_TOKEN}', '<image>')
            #     if verbose:
            #         print(query_to_print, response)
            #     return response, visualization_frames
            ############################################################################################################
        else:
            generation_output = self.generate(
                pixel_values=pixel_values,
                input_ids=input_ids,
                attention_mask=attention_mask,
                attention_visualize=attention_visualize,
                **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(f'{IMG_START_TOKEN}{IMG_END_TOKEN}', '<image>')
                if verbose:
                    print(query_to_print, response)
                return response
        
    def visualize_attention(self, attention_tensor,layer=0, head=None, start_img_token_index=0, end_img_token_index=0, target_aspect_ratio=(0,0)):
        """Vẽ heatmap của attention scores từ layer được chọn và có thể chọn head cụ thể hoặc trung bình."""
        selected_layer = attention_tensor[layer]  # Chọn layer cụ thể
        if head is None:
            averaged_attention = selected_layer.mean(axis=1).squeeze()  # Trung bình qua 14 head
        else:
            averaged_attention = selected_layer[:, head, :, :].squeeze()  # Chọn head cụ thể
        averaged_attention =  np.power(averaged_attention, 0.9)
        heat_maps = []
        for i in range(len(averaged_attention)):  # Duyệt qua 3 beam
            h_target_aspect_ratio = target_aspect_ratio[1]
            if h_target_aspect_ratio == 0 :
                h_target_aspect_ratio = 1
            w_target_aspect_ratio = target_aspect_ratio[0]
            if w_target_aspect_ratio == 0 :
                w_target_aspect_ratio = 1
            img_atten_score = averaged_attention[i].reshape(-1)[start_img_token_index:end_img_token_index]
            img_atten_score = img_atten_score.reshape(h_target_aspect_ratio,w_target_aspect_ratio,16,16)
            img_atten_score = np.transpose(img_atten_score, (0, 2, 1, 3)).reshape(h_target_aspect_ratio*16,w_target_aspect_ratio*16)
            heat_maps.append(img_atten_score)
        return heat_maps



    
    @torch.no_grad()
    def generate(
            self,
            pixel_values: 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,
            attention_visualize: Optional[bool] = False,
            **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 attention_visualize:
            output_attentions = True
            return_dict_in_generate = True
        else:
            output_attentions = False
            return_dict_in_generate = False   
            
        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,
            output_attentions=output_attentions,
            return_dict_in_generate=return_dict_in_generate,
            **generate_kwargs,
        )

        return outputs