# Copyright (c) ModelScope Contributors. All rights reserved.
import os
import safetensors.torch
import torch
from peft import LoraConfig, PeftModel, get_peft_model
from transformers.integrations import is_deepspeed_zero3_enabled
from typing import TYPE_CHECKING, Optional

from swift.utils import deep_getattr, get_logger, get_multimodal_target_regex
from .base import Tuner

logger = get_logger()

if TYPE_CHECKING:
    from swift.arguments import SftArguments


def is_vit_aligner_param(model_arch, parameter_name: str) -> bool:
    for module_prefix in model_arch.vision_tower + model_arch.aligner:
        if f'.{module_prefix}.' in parameter_name:
            return True
    return False


class LoRALLMTuner(Tuner):
    """Full-parameter training of ViT/Aligner while LoRA training LLM"""

    @staticmethod
    def from_pretrained(model: torch.nn.Module, model_id: str, **kwargs) -> torch.nn.Module:
        model = PeftModel.from_pretrained(model, model_id, **kwargs)
        state_dict = safetensors.torch.load_file(os.path.join(model_id, 'vit.safetensors'))
        if is_deepspeed_zero3_enabled():
            import deepspeed
            params_dict = dict(model.named_parameters())
            params_to_load = {name: params_dict[name] for name in state_dict if name in params_dict}
            if params_to_load:
                with deepspeed.zero.GatheredParameters(list(params_to_load.values()), modifier_rank=0):
                    if deepspeed.comm.get_rank() == 0:
                        for name, param in params_to_load.items():
                            param.data.copy_(state_dict[name])
        else:
            model.load_state_dict(state_dict, strict=False)
        model_arch = model.model_meta.model_arch
        for module_prefix in model_arch.vision_tower + model_arch.aligner:
            deep_getattr(model, module_prefix).requires_grad_(True)
        return model

    @staticmethod
    def save_pretrained(
        model: torch.nn.Module,
        save_directory: str,
        state_dict: Optional[dict] = None,
        safe_serialization: bool = True,
        **kwargs,
    ) -> None:
        if state_dict is None:
            state_dict = {}
            for n, p in model.named_parameters():
                if p.requires_grad:
                    state_dict[n] = p.detach().cpu()
        model.save_pretrained(save_directory, state_dict=state_dict, safe_serialization=safe_serialization, **kwargs)
        # vit/aligner
        model_arch = model.model_meta.model_arch
        state_dict = {k: v for k, v in state_dict.items() if is_vit_aligner_param(model_arch, k)}
        safetensors.torch.save_file(
            state_dict, os.path.join(save_directory, 'vit.safetensors'), metadata={'format': 'pt'})

    @staticmethod
    def prepare_model(args: 'SftArguments', model: torch.nn.Module) -> torch.nn.Module:
        model_arch = model.model_meta.model_arch
        target_regex = get_multimodal_target_regex(model)
        logger.info(f'target_regex: {target_regex}')
        lora_config = LoraConfig(
            task_type=args.task_type.upper(), r=args.lora_rank, lora_alpha=args.lora_alpha, target_modules=target_regex)
        model = get_peft_model(model, lora_config)
        for module_prefix in model_arch.vision_tower + model_arch.aligner:
            deep_getattr(model, module_prefix).requires_grad_(True)
        return model