# Copyright (c) ModelScope Contributors. All rights reserved.
import inspect
import torch
import torch.distributed as dist
import transformers
from packaging import version
from PIL import Image
from transformers import PreTrainedModel
from types import MethodType

from swift.template import TemplateType
from swift.utils import is_deepspeed_enabled, to_device
from ..constant import LLMModelType, MLLMModelType
from ..model_arch import ModelArch
from ..model_meta import Model, ModelGroup, ModelMeta
from ..patcher import patch_output_to_input_device
from ..register import ModelLoader, SentenceTransformersLoader, register_model

transformers_5_9 = version.parse(transformers.__version__) >= version.parse('5.9')


class PaligemmaVisionLoader(ModelLoader):

    def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel:
        from transformers import PaliGemmaForConditionalGeneration
        self.auto_model_cls = self.auto_model_cls or PaliGemmaForConditionalGeneration
        return super().get_model(model_dir, *args, **kwargs)


register_model(
    ModelMeta(
        MLLMModelType.paligemma,
        [
            ModelGroup([
                Model('AI-ModelScope/paligemma-3b-pt-224', 'google/paligemma-3b-pt-224'),
                Model('AI-ModelScope/paligemma-3b-pt-448', 'google/paligemma-3b-pt-448'),
                Model('AI-ModelScope/paligemma-3b-pt-896', 'google/paligemma-3b-pt-896'),
            ]),
            ModelGroup([
                Model('AI-ModelScope/paligemma-3b-mix-224', 'google/paligemma-3b-mix-224'),
                Model('AI-ModelScope/paligemma-3b-mix-448', 'google/paligemma-3b-mix-448'),
            ]),
            ModelGroup([
                Model('AI-ModelScope/paligemma2-3b-pt-224', 'google/paligemma2-3b-pt-224'),
                Model('AI-ModelScope/paligemma2-3b-pt-448', 'google/paligemma2-3b-pt-448'),
                Model('AI-ModelScope/paligemma2-3b-pt-896', 'google/paligemma2-3b-pt-896'),
                Model('AI-ModelScope/paligemma2-10b-pt-224', 'google/paligemma2-10b-pt-224'),
                Model('AI-ModelScope/paligemma2-10b-pt-448', 'google/paligemma2-10b-pt-448'),
                Model('AI-ModelScope/paligemma2-10b-pt-896', 'google/paligemma2-10b-pt-896'),
                Model('AI-ModelScope/paligemma2-28b-pt-224', 'google/paligemma2-28b-pt-224'),
                Model('AI-ModelScope/paligemma2-28b-pt-448', 'google/paligemma2-28b-pt-448'),
                Model('AI-ModelScope/paligemma2-28b-pt-896', 'google/paligemma2-28b-pt-896'),
            ]),
            ModelGroup([
                Model('AI-ModelScope/paligemma2-3b-ft-docci-448', 'google/paligemma2-3b-ft-docci-448'),
                Model('AI-ModelScope/paligemma2-10b-ft-docci-448', 'google/paligemma2-10b-ft-docci-448'),
            ]),
        ],
        PaligemmaVisionLoader,
        template=TemplateType.paligemma,
        architectures=['PaliGemmaForConditionalGeneration'],
        model_arch=ModelArch.llava_hf,
        requires=['transformers>=4.41'],
        tags=['vision'],
    ))

register_model(
    ModelMeta(
        LLMModelType.gemma,
        [
            ModelGroup([
                Model('AI-ModelScope/gemma-2b-it', 'google/gemma-2b-it'),
                Model('AI-ModelScope/gemma-2b', 'google/gemma-2b'),
                Model('AI-ModelScope/gemma-7b', 'google/gemma-7b'),
                Model('AI-ModelScope/gemma-7b-it', 'google/gemma-7b-it'),
            ], ),
        ],
        template=TemplateType.gemma,
        architectures=['GemmaForCausalLM'],
        model_arch=ModelArch.llama,
        requires=['transformers>=4.38'],
    ))

register_model(
    ModelMeta(
        LLMModelType.gemma2,
        [
            ModelGroup([
                Model('LLM-Research/gemma-2-2b-it', 'google/gemma-2-2b-it'),
                Model('LLM-Research/gemma-2-2b', 'google/gemma-2-2b'),
                Model('LLM-Research/gemma-2-9b', 'google/gemma-2-9b'),
                Model('LLM-Research/gemma-2-9b-it', 'google/gemma-2-9b-it'),
                Model('LLM-Research/gemma-2-27b', 'google/gemma-2-27b'),
                Model('LLM-Research/gemma-2-27b-it', 'google/gemma-2-27b-it'),
            ], ),
        ],
        template=TemplateType.gemma,
        architectures=['Gemma2ForCausalLM'],
        model_arch=ModelArch.llama,
        requires=['transformers>=4.42'],
    ))


class Gemma3TextLoader(ModelLoader):

    def get_config(self, model_dir):
        # It is strongly recommended to train Gemma3 models with the `eager` attention implementation instead of `sdpa`.
        self.attn_impl = self.attn_impl or 'eager'
        return super().get_config(model_dir)


register_model(
    ModelMeta(
        LLMModelType.gemma3_text,
        [
            ModelGroup([
                Model('LLM-Research/gemma-3-1b-pt', 'google/gemma-3-1b-pt'),
                Model('LLM-Research/gemma-3-1b-it', 'google/gemma-3-1b-it'),
                Model('google/gemma-3-270m', 'google/gemma-3-270m'),
                Model('google/gemma-3-270m-it', 'google/gemma-3-270m-it'),
                Model('google/medgemma-27b-text-it', 'google/medgemma-27b-text-it'),
            ], ),
        ],
        Gemma3TextLoader,
        template=TemplateType.gemma3_text,
        architectures=['Gemma3ForCausalLM'],
        model_arch=ModelArch.llama,
        requires=['transformers>=4.49'],
    ))


class Gemma3VisionLoader(ModelLoader):

    def get_config(self, model_dir):
        # It is strongly recommended to train Gemma3 models with the `eager` attention implementation instead of `sdpa`.
        self.attn_impl = self.attn_impl or 'eager'
        return super().get_config(model_dir)

    def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel:
        from transformers import Gemma3ForConditionalGeneration
        self.auto_model_cls = self.auto_model_cls or Gemma3ForConditionalGeneration
        return super().get_model(model_dir, *args, **kwargs)


register_model(
    ModelMeta(
        MLLMModelType.gemma3_vision,
        [
            ModelGroup([
                Model('LLM-Research/gemma-3-4b-pt', 'google/gemma-3-4b-pt'),
                Model('LLM-Research/gemma-3-4b-it', 'google/gemma-3-4b-it'),
                Model('LLM-Research/gemma-3-12b-pt', 'google/gemma-3-12b-pt'),
                Model('LLM-Research/gemma-3-12b-it', 'google/gemma-3-12b-it'),
                Model('LLM-Research/gemma-3-27b-pt', 'google/gemma-3-27b-pt'),
                Model('LLM-Research/gemma-3-27b-it', 'google/gemma-3-27b-it'),
                Model('google/medgemma-4b-pt', 'google/medgemma-4b-pt'),
                Model('google/medgemma-4b-it', 'google/medgemma-4b-it'),
                Model('google/medgemma-27b-it', 'google/medgemma-27b-it'),
            ], ),
        ],
        Gemma3VisionLoader,
        template=TemplateType.gemma3_vision,
        architectures=['Gemma3ForConditionalGeneration'],
        model_arch=ModelArch.llava_hf,
        requires=['transformers>=4.49'],
    ))


class Gemma3nLoader(ModelLoader):

    def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel:
        from transformers import Gemma3nForConditionalGeneration
        self.auto_model_cls = self.auto_model_cls or Gemma3nForConditionalGeneration
        model = super().get_model(model_dir, *args, **kwargs)
        patch_output_to_input_device(model.model.embed_vision)
        patch_output_to_input_device(model.model.embed_audio)
        return model


register_model(
    ModelMeta(
        MLLMModelType.gemma3n,
        [
            ModelGroup([
                Model('google/gemma-3n-E2B', 'google/gemma-3n-E2B'),
                Model('google/gemma-3n-E4B', 'google/gemma-3n-E4B'),
                Model('google/gemma-3n-E2B-it', 'google/gemma-3n-E2B-it'),
                Model('google/gemma-3n-E4B-it', 'google/gemma-3n-E4B-it'),
            ], ),
        ],
        Gemma3nLoader,
        template=TemplateType.gemma3n,
        architectures=['Gemma3nForConditionalGeneration'],
        model_arch=ModelArch.gemma3n,
        requires=['transformers>=4.53.1'],
    ))

register_model(
    ModelMeta(
        LLMModelType.gemma_emb,
        [
            ModelGroup([
                Model('google/embeddinggemma-300m', 'google/embeddinggemma-300m'),
            ], ),
        ],
        SentenceTransformersLoader,
        template=TemplateType.dummy,
        architectures=['Gemma3TextModel'],
    ))


def _patch_gemma4_forward(model, processor, is_gemma4_unified: bool = False):
    if is_gemma4_unified:
        from transformers.models.gemma4_unified.modeling_gemma4_unified import \
            Gemma4UnifiedModelOutputWithPast as Gemma4ModelOutputWithPast
        from transformers.models.gemma4_unified.modeling_gemma4_unified import (create_masks_for_generate,
                                                                                torch_compilable_check)
    else:
        from transformers.models.gemma4.modeling_gemma4 import (Gemma4ModelOutputWithPast, create_masks_for_generate,
                                                                torch_compilable_check)
    if hasattr(model, 'origin_forward'):
        return

    def _forward_dummy_image(model, inputs_embeds):
        images = [Image.new('RGB', (32, 32), (0, 0, 0))]
        image_inputs = processor.image_processor(images=images, return_tensors='pt')
        image_inputs = to_device(image_inputs, inputs_embeds.device)
        dummy_pixel = image_inputs['pixel_values'].to(model.dtype)
        dummy_pos_ids = image_inputs.get('image_position_ids')
        image_features = model.get_image_features(dummy_pixel, dummy_pos_ids, return_dict=True).pooler_output
        inputs_embeds = inputs_embeds + image_features.mean() * 0.
        return inputs_embeds

    # transformers 5.6.2
    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        pixel_values: torch.FloatTensor | None = None,
        pixel_values_videos: torch.FloatTensor | None = None,
        input_features: torch.FloatTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        input_features_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values=None,
        mm_token_type_ids: torch.LongTensor | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        use_cache: bool | None = None,
        image_position_ids: torch.LongTensor | None = None,
        video_position_ids: torch.LongTensor | None = None,
        per_layer_inputs: torch.Tensor | None = None,
        **kwargs,
    ) -> Gemma4ModelOutputWithPast:
        r"""
        input_features_mask (`torch.FloatTensor]` of shape `(num_images, seq_length)`):
            The attention mask for the input audio.
        image_position_ids (`torch.LongTensor` of shape `(batch_size, max_patches, 2)`, *optional*):
            2D patch position coordinates from the image processor, with `(-1, -1)` indicating padding.
            Passed through to the vision encoder for positional embedding computation.
        video_position_ids (`torch.LongTensor` of shape `(num_videos, num_frames, max_patches, 2)`, *optional*):
            2D patch position coordinates from the video processor, with `(-1, -1)` indicating padding.
            Passed through to the vision encoder for positional embedding computation.
        """
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError('You must specify exactly one of input_ids or inputs_embeds')

        image_mask, video_mask, audio_mask = self.get_placeholder_mask(input_ids, inputs_embeds)
        multimodal_mask = image_mask | video_mask | audio_mask

        # Replace image id with PAD if the image token if OOV, to avoid index-errors
        llm_input_ids = None
        if inputs_embeds is None:
            llm_input_ids = input_ids.clone()
            llm_input_ids[multimodal_mask] = self.config.text_config.pad_token_id
            inputs_embeds = self.get_input_embeddings()(llm_input_ids)

        if per_layer_inputs is None and self.config.get_text_config().hidden_size_per_layer_input:
            pad_embedding = self.language_model.embed_tokens.weight[self.config.text_config.pad_token_id, :]
            pad_embedding = pad_embedding.to(device=multimodal_mask.device)
            llm_inputs_embeds = torch.where(multimodal_mask[..., None], pad_embedding.view(1, 1, -1), inputs_embeds)
            per_layer_inputs = self.language_model.get_per_layer_inputs(llm_input_ids, llm_inputs_embeds)
        else:
            per_layer_inputs = None

        state = input_ids.new_tensor(
            [pixel_values is not None or pixel_values_videos is not None, input_features is not None], dtype=torch.bool)
        if dist.is_initialized() and is_deepspeed_enabled():
            dist.all_reduce(state, dist.ReduceOp.MAX)
        has_image, has_audio = state.tolist()

        # Mixed modality training with both images and videos is not currently supported.
        if pixel_values is None and pixel_values_videos is None and has_image:
            inputs_embeds = _forward_dummy_image(self, inputs_embeds)

        # Merge text and images
        if pixel_values is not None:
            image_features = self.get_image_features(pixel_values, image_position_ids, return_dict=True).pooler_output
            image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)

            # Confirm the number of soft tokens from the vision tower matches the number of slots in the embeddings.
            n_image_tokens = image_mask.sum()
            image_mask = image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
            torch_compilable_check(
                inputs_embeds[image_mask].numel() == image_features.numel(),
                f'Image features and image tokens do not match, tokens: {n_image_tokens}, features:'
                f' {image_features.shape[0]}',
            )

            inputs_embeds = inputs_embeds.masked_scatter(
                image_mask.to(inputs_embeds.device), image_features.to(inputs_embeds.device))

        if pixel_values_videos is not None:
            video_features = self.get_video_features(
                pixel_values_videos, video_position_ids, return_dict=True).pooler_output
            video_features = video_features.to(inputs_embeds.device, inputs_embeds.dtype)

            # Confirm the number of soft tokens from the vision tower matches the number of slots in the embeddings.
            n_video_tokens = video_mask.sum()
            video_mask = video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
            torch_compilable_check(
                inputs_embeds[video_mask].numel() == video_features.numel(),
                f'Video features and video tokens do not match, tokens: {n_video_tokens}, features:'
                f' {video_features.shape[0]}',
            )

            inputs_embeds = inputs_embeds.masked_scatter(
                video_mask.to(inputs_embeds.device), video_features.to(inputs_embeds.device))

        # Merge text and audio
        if input_features is not None and input_features_mask is not None:
            audio_output = self.get_audio_features(input_features, input_features_mask, return_dict=True)
            audio_features = audio_output.pooler_output
            audio_mask_from_encoder = audio_output.attention_mask  # True = valid

            # Strip padding tokens: only keep real (non-padding) audio soft tokens.
            # audio_mask_from_encoder is True for valid positions, False for padding tokens.
            # This mirrors the vision encoder's padding stripping (see Gemma4VisionEncoder.forward).
            audio_features = audio_features[audio_mask_from_encoder]

            n_audio_tokens = audio_mask.sum()
            audio_mask = audio_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
            torch_compilable_check(
                inputs_embeds[audio_mask].numel() == audio_features.numel(),
                f'Audio features and audio tokens do not match, tokens: {n_audio_tokens}, features:'
                f' {audio_features.shape[0] * audio_features.shape[1]}',
            )

            inputs_embeds = inputs_embeds.masked_scatter(
                audio_mask.to(inputs_embeds.device), audio_features.to(inputs_embeds.device))
        elif has_audio and self.audio_tower is not None:
            feature_size = processor.feature_extractor.feature_size
            dummy_features = input_ids.new_zeros([1, 128, feature_size], dtype=self.audio_tower.dtype)
            dummy_mask = input_ids.new_ones([1, 128], dtype=torch.bool)
            audio_output = self.get_audio_features(dummy_features, dummy_mask, return_dict=True)
            audio_features = audio_output.pooler_output
            inputs_embeds = inputs_embeds + audio_features.mean() * 0.

        # It may already have been prepared by, e.g., `generate`
        if position_ids is None:
            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
            position_ids = position_ids.unsqueeze(0)

        bi_vision_attn = self.config.get_text_config().use_bidirectional_attention == 'vision'
        if not isinstance(causal_mask_mapping := attention_mask, dict):
            if bi_vision_attn and not transformers_5_9:
                from transformers.models.gemma4.modeling_gemma4 import create_causal_mask_mapping

                # Larger Gemma 4 models use Gemma 3's bidirectional attention mask for vision inputs
                causal_mask_mapping = create_causal_mask_mapping(
                    self.config,
                    inputs_embeds=inputs_embeds,
                    attention_mask=attention_mask,
                    past_key_values=past_key_values,
                    position_ids=position_ids,
                    mm_token_type_ids=mm_token_type_ids,
                )
            else:
                mask_kwargs = {
                    'config': self.config,
                    'inputs_embeds': inputs_embeds,
                    'attention_mask': attention_mask,
                    'past_key_values': past_key_values,
                    'position_ids': position_ids,
                }
                if bi_vision_attn:
                    from transformers.models.gemma4.modeling_gemma4 import get_block_sequence_ids_for_mask
                    block_sequence_ids = torch.full([*inputs_embeds.size()[:-1]], -1, device=inputs_embeds.device)
                    if mm_token_type_ids is not None:
                        kwargs = {
                            'device': inputs_embeds.device
                        } if 'device' in inspect.signature(get_block_sequence_ids_for_mask).parameters else {}
                        block_sequence_ids = get_block_sequence_ids_for_mask(mm_token_type_ids, **kwargs)

                    mask_kwargs['block_sequence_ids'] = block_sequence_ids

                causal_mask_mapping = create_masks_for_generate(**mask_kwargs)
        kwargs.pop('return_dict', None)
        outputs = self.language_model(
            per_layer_inputs=per_layer_inputs,
            attention_mask=causal_mask_mapping,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            return_dict=True,
            **kwargs,
        )

        return Gemma4ModelOutputWithPast(
            last_hidden_state=outputs.last_hidden_state,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            image_hidden_states=image_features if pixel_values is not None else None,
            audio_hidden_states=audio_features if input_features is not None else None,
        )

    model.origin_forward = model.forward
    model.forward = MethodType(forward, model)


class Gemma4Loader(ModelLoader):

    def get_model(self, model_dir: str, config, processor, model_kwargs) -> PreTrainedModel:
        from transformers import Gemma4ForConditionalGeneration
        self.auto_model_cls = self.auto_model_cls or Gemma4ForConditionalGeneration
        model = super().get_model(model_dir, config, processor, model_kwargs)
        _patch_gemma4_forward(model.model, processor)
        return model


register_model(
    ModelMeta(
        MLLMModelType.gemma4,
        [
            ModelGroup([
                Model('google/gemma-4-E2B', 'google/gemma-4-E2B'),
                Model('google/gemma-4-E2B-it', 'google/gemma-4-E2B-it'),
                Model('google/gemma-4-E4B', 'google/gemma-4-E4B'),
                Model('google/gemma-4-E4B-it', 'google/gemma-4-E4B-it'),
            ],
                       template=TemplateType.gemma4_nothinking),
            ModelGroup([
                Model('google/gemma-4-31B', 'google/gemma-4-31B'),
                Model('google/gemma-4-31B-it', 'google/gemma-4-31B-it'),
                Model('google/gemma-4-26B-A4B', 'google/gemma-4-26B-A4B'),
                Model('google/gemma-4-26B-A4B-it', 'google/gemma-4-26B-A4B-it'),
            ],
                       template=TemplateType.gemma4),
        ],
        Gemma4Loader,
        architectures=['Gemma4ForConditionalGeneration'],
        model_arch=ModelArch.gemma3n,
    ))


class Gemma4UnifiedLoader(ModelLoader):

    def get_model(self, model_dir: str, config, processor, model_kwargs) -> PreTrainedModel:
        from transformers import Gemma4UnifiedForConditionalGeneration
        self.auto_model_cls = self.auto_model_cls or Gemma4UnifiedForConditionalGeneration
        model = super().get_model(model_dir, config, processor, model_kwargs)
        _patch_gemma4_forward(model.model, processor, is_gemma4_unified=True)
        return model


register_model(
    ModelMeta(
        MLLMModelType.gemma4_unified,
        [
            ModelGroup([
                Model('google/gemma-4-12B', 'google/gemma-4-12B'),
                Model('google/gemma-4-12B-it', 'google/gemma-4-12B-it'),
            ],
                       template=TemplateType.gemma4),
        ],
        Gemma4UnifiedLoader,
        architectures=['Gemma4UnifiedForConditionalGeneration'],
        model_arch=ModelArch.gemma4_unified,
        requires=['transformers>=5.10.1'],
    ))


class DiffusionGemmaLoader(ModelLoader):

    def get_model(self, model_dir: str, config, processor, model_kwargs) -> PreTrainedModel:
        from transformers import DiffusionGemmaForBlockDiffusion
        self.auto_model_cls = self.auto_model_cls or DiffusionGemmaForBlockDiffusion
        model = super().get_model(model_dir, config, processor, model_kwargs)
        model.prepare_inputs_for_generation = None
        model.config.use_cache = True
        return model


register_model(
    ModelMeta(
        MLLMModelType.diffusion_gemma,
        [
            ModelGroup([
                Model('google/diffusiongemma-26B-A4B-it', 'google/diffusiongemma-26B-A4B-it'),
            ],
                       template=TemplateType.diffusion_gemma),
        ],
        DiffusionGemmaLoader,
        architectures=['DiffusionGemmaForBlockDiffusion'],
        model_arch=ModelArch.diffusion_gemma,
        requires=['transformers>=5.11'],
    ))
