# Copyright (c) ModelScope Contributors. All rights reserved.
import torch
from functools import partial
from torch import nn
from typing import Any, Dict, List, Literal

from swift.utils import get_env_args, is_deepspeed_enabled
from ..base import Template
from ..constant import MLLMTemplateType
from ..register import register_template
from ..template_inputs import StdTemplateInputs
from ..utils import Context, findall
from ..vision_utils import load_video_internvl, transform_image
from .llm import GptOssTemplateMeta, GptTemplate
from .microsoft import Phi3TemplateMeta
from .utils import ChatmlTemplateMeta


class InternvlTemplate(Template):
    skip_prompt = False
    num_image_token = None
    placeholder_tokens = ['<IMG_CONTEXT>']
    support_padding_free = True

    def init_env_args(self):
        super().init_env_args()
        self.input_size = get_env_args('input_size', int, 448)
        self.max_num = get_env_args('max_num', int, 12)

    def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
                    inputs: StdTemplateInputs) -> List[Context]:
        if self.mode == 'vllm':
            image_context = ['<image>\n']
        else:
            image_context = ['<img>', [-100], '</img>\n']
        return image_context

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        encoded = super()._encode(inputs)
        input_ids = encoded['input_ids']
        idx_list = findall(input_ids, -100)
        pixel_values = None
        images = inputs.images
        if images:
            labels = encoded.get('labels')
            if self.num_image_token is None:
                self.num_image_token = int((self.input_size // 14)**2 * (0.5**2))
            pixel_values_images = [transform_image(image, self.input_size, self.max_num) for image in images]
            pixel_values = torch.cat(pixel_values_images, dim=0).to(self.model_info.torch_dtype)
            image_bs = pixel_values.shape[0]

            idx, idx2 = idx_list[0], idx_list[-1]  # remove [-100, -100]
            img_tokens: List[int] = self.processor.encode(
                '<IMG_CONTEXT>', add_special_tokens=False) * self.num_image_token * image_bs
            input_ids = input_ids[:idx] + img_tokens + input_ids[idx2 + 1:]
            if labels is not None:
                labels = labels[:idx] + [-100] * len(img_tokens) + labels[idx2 + 1:]
            encoded['input_ids'] = input_ids
            encoded['labels'] = labels
        encoded['pixel_values'] = pixel_values
        return encoded

    def forward_context(self, model, inputs):
        model_name = model.language_model.__class__.__name__.lower()
        if self.padding_free and 'internlm2' in model_name:
            position_ids = inputs['position_ids']
            modeling_module = model.language_model.model.layers[0].attention.__class__
            return self._patch_flash_attention_forward(modeling_module, position_ids, use_new_func=True)
        else:
            return super().forward_context(model, inputs)

    def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]:
        embedding = model.get_input_embeddings()
        device = embedding.weight.device
        input_ids = inputs['input_ids']
        inputs_embeds = embedding(input_ids).to(device=device)
        pixel_values = inputs.get('pixel_values')
        if pixel_values is not None:
            pixel_values = pixel_values.to(device=device)
            vit_embeds = model.extract_feature(pixel_values).to(device=device)
            selected = (input_ids == self.processor.encode('<IMG_CONTEXT>', add_special_tokens=False)[0])
            inputs_embeds[selected] = vit_embeds.reshape(-1, vit_embeds.shape[-1]).to(dtype=inputs_embeds.dtype)
        elif is_deepspeed_enabled():
            dummy_pixel_values = torch.zeros((1, 3, 32, 32), device=device, dtype=inputs_embeds.dtype)
            vit_embeds = model.extract_feature(dummy_pixel_values).to(device=device)
            inputs_embeds += vit_embeds.mean() * 0.
        return {'inputs_embeds': inputs_embeds}


register_template(
    ChatmlTemplateMeta(
        MLLMTemplateType.internvl,
        default_system='You are an AI assistant whose name is InternLM (书生·浦语).',
        template_cls=InternvlTemplate,
        auto_add_bos=True))
register_template(
    Phi3TemplateMeta(
        MLLMTemplateType.internvl_phi3,
        default_system='You are an AI assistant whose name is Phi-3.',
        template_cls=InternvlTemplate,
        auto_add_bos=True))


class Internvl2Template(InternvlTemplate):
    VIDEO_SEGMENTS = 8

    def init_env_args(self):
        super().init_env_args()
        self.video_max_num = get_env_args('video_max_num', int, 1)
        self.video_segments = get_env_args('video_segments', int, self.VIDEO_SEGMENTS)

    def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
                    inputs: StdTemplateInputs) -> List[Context]:
        image_context = super().replace_tag('image', index, inputs)
        if media_type == 'image':
            return image_context
        elif media_type == 'video':
            load_video = partial(load_video_internvl, num_segments=self.video_segments)
            return self.replace_video2image(load_video, inputs, lambda i: [f'Frame{i + 1}: '] + image_context)

    def replace_ref(self, ref: str, index: int, inputs: StdTemplateInputs) -> List[Context]:
        return [f'<ref>{ref}</ref>']

    def replace_bbox(self, bbox: List[int], index: int, inputs: StdTemplateInputs) -> List[Context]:
        return [f'<box>[{bbox}]</box>']

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        encoded = super(InternvlTemplate, self)._encode(inputs)
        input_ids = encoded['input_ids']
        idx_list = findall(input_ids, -100)
        labels = encoded['labels']
        loss_scale = encoded.get('loss_scale', None)
        images = inputs.images
        if images:
            has_video = bool(inputs.videos)
            if self.num_image_token is None:
                self.num_image_token = int((self.input_size // 14)**2 * (0.5**2))
            max_num = self.max_num
            if has_video:
                max_num = self.video_max_num
            pixel_values = [transform_image(image, self.input_size, max_num) for image in images]
            num_patches = [pv.shape[0] for pv in pixel_values]
            pixel_values = torch.cat(pixel_values).to(self.model_info.torch_dtype)
        else:
            pixel_values = None
            num_patches = []
        assert len(num_patches) == len(
            idx_list), f'len(num_patches): {len(num_patches)}, len(idx_list): {len(idx_list)}'

        def _get_new_tokens(i):
            img_tokens: List[int] = self.processor.encode(
                '<IMG_CONTEXT>', add_special_tokens=False) * self.num_image_token * num_patches[i]
            return img_tokens

        encoded['input_ids'], encoded['labels'], encoded['loss_scale'] = self._extend_tokens(
            input_ids, labels, loss_scale, idx_list, _get_new_tokens)
        encoded['pixel_values'] = pixel_values
        return encoded


_internvl2_system = '你是由上海人工智能实验室联合商汤科技开发的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。'
register_template(
    ChatmlTemplateMeta(
        MLLMTemplateType.internvl2,
        default_system=_internvl2_system,
        template_cls=Internvl2Template,
    ))

register_template(
    Phi3TemplateMeta(
        MLLMTemplateType.internvl2_phi3,
        default_system=_internvl2_system,
        template_cls=Internvl2Template,
    ))

register_template(
    ChatmlTemplateMeta(
        MLLMTemplateType.internvl2_5,
        template_cls=Internvl2Template,
        default_system='你是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。'))

register_template(ChatmlTemplateMeta(MLLMTemplateType.internvl3_5, template_cls=Internvl2Template))


class Internvl3_5GPTTemplate(Internvl2Template, GptTemplate):
    pass


register_template(GptOssTemplateMeta(MLLMTemplateType.internvl3_5_gpt, template_cls=Internvl3_5GPTTemplate))


class InternvlhfTemplate(Internvl2Template):

    def init_env_args(self):
        Template.init_env_args(self)

    def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
                    inputs: StdTemplateInputs) -> List[Context]:
        assert media_type in ['image', 'video']
        if media_type == 'video':
            if self.mode == 'vllm':
                return Template.replace_tag(self, 'video', index, inputs)
            else:
                return [[-200]]
        else:
            if self.mode == 'vllm':
                return ['<IMG_CONTEXT>']
            else:
                return ['<img>', [-100], '</img>\n']

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        import numpy as np
        from transformers.image_utils import concatenate_list, make_flat_list_of_images
        from transformers.video_utils import make_batched_videos

        from swift.template.vision_utils import load_video_hf
        encoded = super(InternvlTemplate, self)._encode(inputs)
        input_ids = encoded['input_ids']
        labels = encoded['labels']
        loss_scale = encoded.get('loss_scale', None)
        images = inputs.images
        videos = inputs.videos
        image_num_patches_indices = np.array([0])
        video_num_patches_indices = np.array([0])
        video_patch_indices = np.array([0])
        image_num_patches = []
        video_num_patches = []
        image_video_patches = []
        image_idx_list = []
        video_idx_list = []
        image_pixel_values = None
        video_pixel_values = None

        if images:
            # InternS1Processor
            image_idx_list = findall(input_ids, -100)
            images = make_flat_list_of_images(images)
            image_inputs = self.processor.image_processor(images=images, crop_to_patches=True, return_tensors='pt')
            image_num_patches = image_inputs.pop('num_patches')
            image_pixel_values = image_inputs.pop('pixel_values').to(self.model_info.torch_dtype)
            image_num_patches_indices = np.cumsum(image_num_patches)
        if videos:
            video_idx_list = findall(input_ids, -200)
            videos, _ = load_video_hf(videos)
            videos = make_batched_videos(videos)
            video_inputs = self.processor.video_processor(videos=videos, return_tensors='pt')
            video_pixel_values = video_inputs.pop('pixel_values_videos').to(self.model_info.torch_dtype)
            num_frames_per_video = [len(video) for video in video_pixel_values]
            video_num_patches = [1 for frames in num_frames_per_video for _ in range(frames)]
            video_patch_indices = np.cumsum(num_frames_per_video)
            video_num_patches_indices = np.cumsum(video_num_patches)
            video_pixel_values = video_pixel_values.flatten(0, 1)

        def merge_and_sort(image_idx_list: List[int], video_idx_list: List[int]) -> tuple:
            """Merge and sort image and video index lists while preserving their relative order."""
            merged = []
            is_image_list = []
            i, j = 0, 0

            while i < len(image_idx_list) and j < len(video_idx_list):
                if image_idx_list[i] < video_idx_list[j]:
                    merged.append(image_idx_list[i])
                    i += 1
                    is_image_list.append(True)
                else:
                    merged.append(video_idx_list[j])
                    j += 1
                    is_image_list.append(False)
            # Add remaining elements
            merged.extend(image_idx_list[i:])
            is_image_list.extend([True] * (len(image_idx_list) - i))
            merged.extend(video_idx_list[j:])
            is_image_list.extend([False] * (len(video_idx_list) - j))
            return merged, is_image_list

        # Merge and sort the index lists
        idx_list, is_image_list = merge_and_sort(image_idx_list, video_idx_list)

        # Validate the lengths
        if images and len(image_idx_list) > 0:
            assert len(image_num_patches_indices) == len(image_idx_list)
        if videos and len(video_idx_list) > 0:
            assert len(video_patch_indices) == len(video_idx_list)

        def _get_new_tokens(i):
            if is_image_list[i]:
                # Find the corresponding image index
                image_idx = sum(is_image_list[:i])
                start = image_num_patches_indices[image_idx - 1] if image_idx > 0 else 0
                end = image_num_patches_indices[image_idx]
                image_seq_length = self.processor.image_seq_length
                image_video_patches.append(image_pixel_values[start:end])
                img_tokens: List[int] = self.processor.encode(
                    '<IMG_CONTEXT>', add_special_tokens=False) * image_seq_length * image_num_patches[image_idx]
            else:
                # Find the corresponding video index
                video_idx = i - sum(is_image_list[:i])
                current_patch = video_patch_indices[video_idx - 1] if video_idx > 0 else 0
                end_patch = video_patch_indices[video_idx]

                start = video_num_patches_indices[current_patch] if video_idx > 0 else 0
                end = video_num_patches_indices[end_patch - 1]
                image_video_patches.append(video_pixel_values[start:end])
                image_seq_length = self.processor.image_seq_length
                num_patches = list(video_num_patches[current_patch:end_patch])
                video_prompt = ''.join(
                    f"Frame{i + 1}: <img>{'<IMG_CONTEXT>' * image_seq_length * num_patches[i]}</img>\n"
                    for i in range(len(num_patches)))
                img_tokens = self.processor.encode(video_prompt, add_special_tokens=False)
            return img_tokens

        encoded['input_ids'], encoded['labels'], encoded['loss_scale'] = self._extend_tokens(
            input_ids, labels, loss_scale, idx_list, _get_new_tokens)
        if images or videos:
            encoded['pixel_values'] = concatenate_list(image_video_patches)
        return encoded

    def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]:
        embedding = model.get_input_embeddings()
        device = embedding.weight.device
        input_ids = inputs['input_ids']
        inputs_embeds = embedding(input_ids).to(device=device)
        pixel_values = inputs.get('pixel_values')
        if pixel_values is not None:
            pixel_values = pixel_values.to(device=device)
            image_features = model.model.get_image_features(
                pixel_values,
                vision_feature_layer=self.config.vision_feature_layer,
                vision_feature_select_strategy=self.config.vision_feature_select_strategy,
            )
            special_image_mask = input_ids == self.config.image_token_id
            special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
            image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
        elif is_deepspeed_enabled():
            dummy_pixel_values = torch.zeros((1, 3, 32, 32), device=device, dtype=inputs_embeds.dtype)
            image_features = model.model.get_image_features(
                dummy_pixel_values,
                vision_feature_layer=self.config.vision_feature_layer,
                vision_feature_select_strategy=self.config.vision_feature_select_strategy,
            )
            inputs_embeds = inputs_embeds + image_features.mean() * 0.
        return {'inputs_embeds': inputs_embeds}


INTERNS1_DEFAULT_SYSTEM = ('You are an expert reasoner with extensive experience in all areas. '
                           'You approach problems through systematic thinking and rigorous reasoning. '
                           'Your response should reflect deep understanding and precise logical thinking, '
                           'making your solution path and reasoning clear to others. '
                           'Please put your thinking process within <think>...</think> tags.')

register_template(
    ChatmlTemplateMeta(
        MLLMTemplateType.interns1,
        template_cls=InternvlhfTemplate,
        default_system=INTERNS1_DEFAULT_SYSTEM,
        is_thinking=True,
        thinking_prefix='<think>',
    ))

register_template(ChatmlTemplateMeta(MLLMTemplateType.internvl_hf, template_cls=InternvlhfTemplate))
