# Copyright (c) ModelScope Contributors. All rights reserved.
import re
import torch
import torch.distributed as dist
from megatron.core import mpu
from megatron.core.extensions.transformer_engine import TEGroupedLinear, TELayerNormColumnParallelLinear, TELinear
from megatron.core.inference.communication_utils import recv_from_prev_pipeline_rank_, send_to_next_pipeline_rank
from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
from megatron.core.packed_seq_params import PackedSeqParams
from megatron.core.ssm.mamba_context_parallel import _undo_attention_load_balancing
from megatron.core.transformer.moe.router import TopKRouter
from torch import nn
from transformers.utils import is_torch_npu_available

from swift.tuners import LoraConfig, Swift
from swift.utils import (activate_parameters, deep_getattr, find_layers, freeze_parameters, get_logger,
                         get_model_parameter_info)
from swift.utils import get_packed_seq_params as _get_packed_seq_params

logger = get_logger()


def find_all_linears(model, extra_layers=None):

    def _cond(name, module):
        if (extra_layers and isinstance(module, tuple(extra_layers))) or name != 'output_layer' and isinstance(
                module, (TELinear, TELayerNormColumnParallelLinear, TEGroupedLinear, nn.Linear)):
            return True
        return False

    return find_layers(model, _cond)


def find_router(model):
    return find_layers(model, lambda name, module: isinstance(module, TopKRouter))


def find_embedding(model):
    return find_layers(model, lambda name, module: isinstance(module, LanguageModelEmbedding))


def get_multimodal_target_regex(
    args,
    model,
    *,
    freeze_llm: bool = False,
    freeze_vit: bool = True,
    freeze_aligner: bool = True,
    include_embedding: bool = False,
    include_router: bool = False,
) -> str:
    megatron_model_meta = args.megatron_model_meta
    modules = []
    visual_cls = megatron_model_meta.visual_cls
    vision_tower = [f'visual.{vit}' for vit in visual_cls._vision_tower]
    aligner = [f'visual.{aligner}' for aligner in visual_cls._aligner]
    if not freeze_llm:
        modules.append('language_model')
    if not freeze_vit:
        modules += vision_tower
    if not freeze_aligner:
        modules += aligner
    assert len(modules) > 0, f'modules: {modules}'
    extra_layers = []
    if include_embedding:
        extra_layers.append(LanguageModelEmbedding)
    if include_router:
        extra_layers.append(TopKRouter)

    res = []
    for module in modules:
        rejected_modules = []
        if not freeze_vit:
            for _aligner in aligner:
                if _aligner.startswith(f'{module}.'):
                    rejected_modules.append(_aligner)

        sub_module = deep_getattr(model, module)
        if sub_module is None:
            continue
        target_modules = find_all_linears(sub_module, extra_layers)
        if not target_modules:
            continue
        target_modules = [tm for tm in target_modules if tm]
        target_pattern = rf'.*\.({"|".join(target_modules)})' if target_modules else ''
        rejected_pattern = rf'(?!({"|".join(rejected_modules)}))' if rejected_modules else ''
        res.append(rf'{rejected_pattern}{re.escape(module)}(?=\.){target_pattern}')

    return rf'^({"|".join(res)})$'


def get_target_modules(args, model):
    if isinstance(args.target_modules, str):
        return args.target_modules
    target_modules = args.target_modules.copy()
    if 'all-linear' in target_modules:
        if args.is_multimodal and not args.language_model_only:
            if args.tuner_type == 'lora_llm':
                kwargs = {
                    'freeze_llm': False,
                    'freeze_vit': True,
                    'freeze_aligner': True,
                }
            else:  # lora
                kwargs = {
                    'freeze_llm': args.freeze_llm,
                    'freeze_vit': args.freeze_vit,
                    'freeze_aligner': args.freeze_aligner,
                }
            return get_multimodal_target_regex(
                args,
                model,
                include_embedding='all-embedding' in target_modules,
                include_router='all-router' in target_modules,
                **kwargs,
            )
        else:
            target_modules.remove('all-linear')
            target_modules += find_all_linears(model)
    if 'all-embedding' in target_modules:
        target_modules.remove('all-embedding')
        target_modules += find_embedding(model)
    if 'all-router' in target_modules:
        target_modules.remove('all-router')
        target_modules += find_router(model)
    return target_modules


def get_modules_to_save(args, model):
    if args.task_type == 'seq_cls':
        args.modules_to_save.append('output_layer')
    modules_to_save = args.modules_to_save.copy()
    if 'all-embedding' in args.modules_to_save:
        modules_to_save.remove('all-embedding')
        modules_to_save += find_embedding(model)
    return modules_to_save


def prepare_adapter(args, model):
    target_modules = get_target_modules(args, model)
    modules_to_save = get_modules_to_save(args, model)
    lora_kwargs = {
        'r': args.lora_rank,
        'target_modules': target_modules,
        'lora_alpha': args.lora_alpha,
        'lora_dropout': args.lora_dropout,
        'bias': args.lora_bias,
        'modules_to_save': modules_to_save,
        'use_rslora': args.use_rslora,
    }
    lora_config = LoraConfig(task_type='CAUSAL_LM', lora_dtype=args.lora_dtype, **lora_kwargs)
    logger.info(f'lora_config: {lora_config}')
    model = Swift.prepare_model(model, lora_config)
    if args.mcore_ref_adapter or args.ref_adapters:
        model.add_adapter('ref_adapter', lora_config)
        model.base_model._cast_adapter_dtype(adapter_name='ref_adapter', autocast_adapter_dtype=True)
        for n, p in model.named_parameters():
            if '.ref_adapter.' in n:
                p.requires_grad = False
    return model


def _prepare_full_vit(args, model):
    megatron_model_meta = args.megatron_model_meta
    visual_cls = megatron_model_meta.visual_cls
    vision_tower = [f'visual.{vit}' for vit in visual_cls._vision_tower]
    aligner = [f'visual.{aligner}' for aligner in visual_cls._aligner]
    for module_prefix in vision_tower + aligner:
        module = deep_getattr(model, module_prefix)
        if module is not None:
            module.requires_grad_(True)


def prepare_mcore_model(args, model):
    if args.tuner_type == 'full':
        freeze_parameters(model, args.freeze_parameters_ratio, args.freeze_parameters, args.freeze_parameters_regex)
        if args.trainable_parameters or args.trainable_parameters_regex:
            activate_parameters(model, args.trainable_parameters, args.trainable_parameters_regex)
    elif args.tuner_type in {'lora', 'lora_llm'}:
        model = prepare_adapter(args, model)
        if args.tuner_type == 'lora_llm':
            _prepare_full_vit(args, model)
    logger.info(f'model: {model}')
    logger.info_if(
        f'[rank{dist.get_rank()}] model_parameter_info: {get_model_parameter_info(model)}',
        cond=mpu.get_data_parallel_rank() == 0)
    return model


def forward_step_helper(model, inputs, dtype=None):
    config = model.config
    dtype = dtype or config.params_dtype
    if not mpu.is_pipeline_first_stage():
        recv_shape_buffer = torch.empty((3, ), device=torch.cuda.current_device(), dtype=torch.int64)
        recv_from_prev_pipeline_rank_(recv_shape_buffer)
        recv_buffer = torch.empty(recv_shape_buffer.tolist(), device=torch.cuda.current_device(), dtype=dtype)
        recv_from_prev_pipeline_rank_(recv_buffer)
        model.set_input_tensor(recv_buffer)
    output_tensor = model(**inputs)
    if not mpu.is_pipeline_last_stage():
        recv_shape_buffer = torch.tensor(output_tensor.shape, device=torch.cuda.current_device(), dtype=torch.int64)
        send_to_next_pipeline_rank(recv_shape_buffer)
        send_to_next_pipeline_rank(output_tensor)
        output_tensor = None

    return output_tensor


def get_padding_to(args):
    padding_to = None
    if args.tensor_model_parallel_size > 1 and args.sequence_parallel:
        padding_to = args.tensor_model_parallel_size
    if args.context_parallel_size > 1:
        padding_to = (padding_to or 1) * args.context_parallel_size
    origin_padding_to = padding_to
    fp8_format = getattr(args, 'fp8_format', None) or getattr(args, 'fp8', None)
    fp4_format = getattr(args, 'fp4_format', None) or getattr(args, 'fp4', None)
    if args.fp8_recipe == 'blockwise':
        padding_to = (padding_to or 1) * 128
    elif args.fp8_recipe == 'mxfp8':
        # MXFP8 uses a block size of 32. Under sequence parallel, the sequence is
        # split across TP ranks, so each per-rank shard (seq_len / TP) must itself
        # be divisible by 32. Pad the total length to TP * 32 to guarantee this.
        padding_to = (padding_to or 1) * 32
    elif fp8_format is not None or fp4_format is not None:
        padding_to = (padding_to or 1) * 16
    if args.attention_backend == 'fused':
        padding_to = max(padding_to or 1, ((origin_padding_to) or 1) * 64)
    return padding_to


def get_packed_seq_params(args, position_ids: torch.Tensor) -> PackedSeqParams:
    params = _get_packed_seq_params(position_ids)
    packed = PackedSeqParams(
        cu_seqlens_q=params['cu_seq_lens_q'],
        cu_seqlens_kv=params['cu_seq_lens_k'],
        max_seqlen_q=params['max_length_q'],
        max_seqlen_kv=params['max_length_k'],
        qkv_format='thd',
    )
    if hasattr(packed, 'cp_partition_mode'):
        packed.cp_partition_mode = args.cp_partition_mode

    if is_torch_npu_available():
        packed.cu_seqlens_q_padded = params['cu_seq_lens_q']
        packed.cu_seqlens_kv_padded = params['cu_seq_lens_k']

    return packed


def reconstruct_tensor_cp(tensor, packed_seq_params, dim=1) -> torch.Tensor:
    """In CP mode, all-gather and undo the load-balanced (zigzag) chunking
    produced by ``split_cp_inputs``, restoring the full sequence in original
    token order along ``dim``.

    Args:
        tensor: CP-sharded local tensor whose sequence dim is at ``dim``.
        packed_seq_params: ``PackedSeqParams`` for THD inputs, or ``None`` for
            regular ``[B, S, ...]`` inputs.
        dim: Sequence dimension index of ``tensor`` (default: 1).

    Returns:
        torch.Tensor: Full-sequence tensor with the same shape as ``tensor``
        except the size at ``dim`` is multiplied by ``cp_size``.
    """

    cp_size = mpu.get_context_parallel_world_size()
    if cp_size <= 1:
        return tensor

    cp_rank = mpu.get_context_parallel_rank()
    cp_group = mpu.get_context_parallel_group()

    # All-gather across CP ranks (preserve local autograd graph for `tensor`).
    output_list = [torch.empty_like(tensor) for _ in range(cp_size)]
    torch.distributed.all_gather(output_list, tensor.contiguous(), group=cp_group)
    output_list[cp_rank] = tensor
    gathered = torch.cat(output_list, dim=dim)

    # `_undo_attention_load_balancing` assumes sequence dim is 0; transpose if needed.
    if dim != 0:
        gathered = gathered.transpose(0, dim).contiguous()
    out = _undo_attention_load_balancing(gathered, cp_size, packed_seq_params)
    if dim != 0:
        out = out.transpose(0, dim).contiguous()
    return out
