# Copyright (c) ModelScope Contributors. All rights reserved.

import math
import torch
import torch.distributed as dist
import torch.nn as nn
from contextlib import contextmanager, nullcontext
from megatron.core import mpu
from megatron.core.extensions.transformer_engine import TEDotProductAttention
from megatron.core.ssm.mamba_context_parallel import _undo_attention_load_balancing
from megatron.core.tensor_parallel import VocabParallelEmbedding
from megatron.core.tensor_parallel.mappings import (gather_from_sequence_parallel_region,
                                                    gather_from_tensor_model_parallel_region)
from typing import Any, Dict

from swift.utils import HfConfigFactory, get_logger, to_device, to_float_dtype
from .megatron_lm_utils import get_batch_on_this_cp_rank
from .utils import forward_step_helper, get_packed_seq_params, get_padding_to

logger = get_logger()


def _test_params_sum(model):
    total_sum = 0
    zero_count = 0
    n_parameter = 0
    for n, p in model.named_parameters():
        n_parameter += 1
        sum_ = p.to(device='cuda', dtype=torch.float32).abs().sum().cpu().item()
        if sum_ == 0 and '.lora_B.' not in n:
            zero_count += 1
            logger.warning(f'n: {n}, sum: {sum_}')
        elif math.isnan(sum_) or math.isinf(sum_) or sum_ > 1e10:
            logger.warning(f'n: {n}, sum: {sum_}')
        else:
            total_sum += sum_
    cond = mpu.get_data_parallel_rank() == 0
    logger.info_if(f'n_parameter: {n_parameter}', cond=cond)
    logger.info_if(f'total_sum: {total_sum}', cond=cond)
    logger.info_if(f'zero_count: {zero_count}', cond=cond)


def _find_modules(model, recurse: bool = True, prefix='', ignore_modules=None):
    ignore_modules = ignore_modules or []
    for k in ignore_modules:
        if prefix.startswith(k):
            return []
    else:
        named_children = list(model.named_children())

    modules = []
    for n, module in named_children:
        if module.__class__ is nn.ModuleList:
            modules += _find_modules(module, False, prefix=f'{prefix}{n}.', ignore_modules=ignore_modules)
        elif recurse:
            modules += _find_modules(module, prefix=f'{prefix}{n}.', ignore_modules=ignore_modules)
        else:
            modules.append(module)
    if not named_children:
        modules.append(model)
    return modules


@contextmanager
def _model_cpu_forward_context(modules,
                               torch_dtype=None,
                               compute_device=None,
                               share_embedding: bool = False,
                               target_device='cpu'):
    for module in modules:
        try:
            origin_torch_dtype = next(module.parameters()).dtype
        except StopIteration:
            pass
        else:
            break
    embeddings = None
    if share_embedding:
        embeddings = [module for module in modules if isinstance(module, (nn.Embedding, VocabParallelEmbedding))]

    def _to_cuda_hook(module, args):
        if compute_device is not None or torch_dtype is not None:
            module.to(device=compute_device, dtype=torch_dtype)
            args = to_float_dtype(args, dtype=torch_dtype)
        return args

    def _to_cpu_hook(module, args, output):
        if share_embedding and module in embeddings or 'rotaryemb' in module.__class__.__name__.lower():
            return
        module.to(device=target_device, dtype=origin_torch_dtype)

    hooks = []
    for module in modules:
        hooks.append(module.register_forward_pre_hook(_to_cuda_hook))
        hooks.append(module.register_forward_hook(_to_cpu_hook))
    try:
        yield
    finally:
        for hook in hooks:
            hook.remove()


def get_examples(mm_type: str) -> Dict[str, Any]:
    if mm_type == 'image':
        data = {
            'messages': [{
                'role': 'user',
                'content': '<image>describe the image.'
            }, {
                'role':
                'assistant',
                'content':
                'The image depicts a close-up of a kitten with striking features. '
                'The kitten has a white and gray coat with distinct black stripes, '
                'particularly noticeable on its face and ears. Its eyes are large '
                'and expressive, with a captivating blue hue that stands out against '
                "the darker fur around them. The kitten's nose is small and pink, "
                'and it has long, delicate whiskers extending from either side of its mouth. '
                "The background is blurred, drawing attention to the kitten's face and "
                'making it the focal point of the image. The overall impression is '
                'one of cuteness and charm.'
            }],
            'images': ['http://modelscope-open.oss-cn-hangzhou.aliyuncs.com/images/cat.png']
        }
    elif mm_type == 'audio':
        data = {
            'messages': [{
                'role': 'user',
                'content': '<audio>Caption the audio.'
            }, {
                'role': 'assistant',
                'content': "The audio contains a male voice speaking the phrase '今天天气真好呀' in Mandarin."
            }],
            'audios': ['http://modelscope-open.oss-cn-hangzhou.aliyuncs.com/images/weather.wav']
        }
    else:  # text
        data = {
            'messages': [
                {
                    'role': 'user',
                    'content': 'Introduction to ms-swift.'
                },
                {
                    'role':
                    'assistant',
                    'content':
                    'ms-swift is an official framework provided by the ModelScope community for fine-tuning '
                    'and deploying large language models and multi-modal large models.'
                },
            ]
        }
    return data


def broadcast_mg_logits(mg_logits=None, src_rank=None):
    if not dist.is_initialized():
        return
    rank = dist.get_rank()
    if src_rank is None:
        src_rank = dist.get_world_size() - 1
    if rank == src_rank:
        meta = [tuple(mg_logits.shape), str(mg_logits.dtype).split('.', 1)[1]]
    else:
        meta = [None, None]

    dist.broadcast_object_list(meta, src=src_rank)
    shape, dtype = meta
    dtype = getattr(torch, dtype)

    if rank != src_rank:
        mg_logits = torch.empty(shape, dtype=dtype, device='cuda')

    dist.broadcast(mg_logits, src=src_rank)

    return mg_logits


@contextmanager
def _patch_attention_fp32(compute_dtype):
    forward = TEDotProductAttention.forward

    def new_forward(self, query_layer, key_layer, value_layer, *args, **kwargs):
        torch_dtype = query_layer.dtype
        query_layer = query_layer.to(compute_dtype)
        key_layer = key_layer.to(compute_dtype)
        value_layer = value_layer.to(compute_dtype)
        res = forward(self, query_layer, key_layer, value_layer, *args, **kwargs)
        res = res.to(dtype=torch_dtype)
        return res

    TEDotProductAttention.forward = new_forward
    try:
        yield
    finally:
        TEDotProductAttention.forward = forward


def test_convert_precision(args, hf_model, mg_model, template, test_convert_dtype=None):
    if test_convert_dtype is None:
        test_convert_dtype = getattr(args, 'test_convert_dtype', torch.float32)
    template.set_mode('train')
    _test_params_sum(mg_model)

    config = mg_model.config
    is_multimodal = config.is_multimodal
    if is_multimodal:
        test_mm_type = getattr(config, 'test_mm_type', 'image')
    else:
        test_mm_type = 'text'
    mg_language_model = mg_model.language_model if is_multimodal else mg_model
    if mg_language_model.config.fp8 is not None:
        raise ValueError('fp8 models currently do not support testing convert_precision. '
                         'Please set `--test_convert_precision false`.')
    share_embedding = mg_language_model.share_embeddings_and_output_weights
    if hf_model is not None:
        hf_model.eval()
        if dist.get_world_size() == 1:
            _test_params_sum(hf_model)
        inputs = template.encode(get_examples(test_mm_type), return_length=True)
        hf_inputs = to_device(template.data_collator([inputs]), 'cuda')
        template.register_post_encode_hook([hf_model])
        HfConfigFactory.set_config_attr(hf_model.config, 'use_cache', False)
        model_arch = hf_model.model_meta.model_arch
        ignore_modules = (model_arch.vision_tower + model_arch.aligner) if is_multimodal else []
        hf_modules = _find_modules(hf_model, ignore_modules=ignore_modules)
        with torch.inference_mode(), _model_cpu_forward_context(
                hf_modules, test_convert_dtype, share_embedding=share_embedding):
            hf_inputs.pop('text_position_ids', None)
            hf_logits = hf_model(**hf_inputs).logits
            hf_logits = hf_logits.to('cuda')
        hf_model.to('cpu')

    template.use_megatron = True
    inputs = [
        template.encode(get_examples(test_mm_type), return_length=True) for _ in range(2 if args.padding_free else 1)
    ]
    mg_inputs = to_device(template.data_collator(inputs, padding_to=get_padding_to(args)), 'cuda')
    mg_model.eval()
    # thd
    text_position_ids = mg_inputs.pop('text_position_ids', None)
    if text_position_ids is None:
        text_position_ids = mg_inputs.get('position_ids')
    if args.padding_free:
        mg_inputs['packed_seq_params'] = get_packed_seq_params(args, text_position_ids)
    mg_language_model.config.fp8 = None  # compat fp8
    mg_modules = _find_modules(mg_language_model, ignore_modules=['visual'])
    for key in ['labels', 'seq_lens', 'attention_mask_2d']:
        mg_inputs.pop(key, None)
    mg_inputs = get_batch_on_this_cp_rank(args, mg_inputs)
    _param = next(mg_language_model.parameters())
    mg_dtype = _param.dtype
    mg_device = _param.device
    if args.model_type == 'minimax_m2':
        # router to bfloat16 (expert_bias). No need to do this when actually training.
        for n, m in mg_language_model.named_modules():
            if n.endswith('router'):
                m.to(mg_dtype)
    if getattr(config, 'enable_hyper_connections', False):
        for param in mg_language_model.decoder.parameters(recurse=False):
            param.data = param.data.cuda()
    attention_context = (
        _patch_attention_fp32(mg_dtype) if args.attention_backend.name in {'flash', 'fused'} else nullcontext())
    with torch.inference_mode(), _model_cpu_forward_context(
            mg_modules, test_convert_dtype, 'cuda', share_embedding=share_embedding,
            target_device=mg_device), attention_context:
        mg_logits = forward_step_helper(mg_model, mg_inputs, dtype=test_convert_dtype)
        if args.tensor_model_parallel_size > 1 and args.task_type != 'seq_cls':
            if mg_logits is not None:
                mg_logits = gather_from_tensor_model_parallel_region(mg_logits)
        if args.context_parallel_size > 1:
            if mg_logits is not None:
                mg_logits = gather_from_sequence_parallel_region(
                    mg_logits.transpose(0, 1), group=mpu.get_context_parallel_group())
                # Contiguous CP already gathers ranks in original token order, so the
                # zigzag un-balancing must be skipped (it is only correct for zigzag split).
                if getattr(args, 'cp_partition_mode', 'zigzag') != 'contiguous':
                    mg_logits = _undo_attention_load_balancing(mg_logits, args.context_parallel_size)
                mg_logits = mg_logits.transpose(0, 1)
    mg_logits = broadcast_mg_logits(mg_logits)
    if hf_model is None:
        return
    if args.task_type == 'seq_cls':
        mg_logits = mg_logits[:, -1]
        mean_diff = (mg_logits - hf_logits).abs().mean().item()
        max_diff = (mg_logits - hf_logits).abs().max().item()
        print(f'mean_diff: {mean_diff}, max_diff: {max_diff}')
    else:
        mg_logits = mg_logits[:, :hf_logits.shape[1]]
        token_mean_diff = (mg_logits - hf_logits).abs().mean(dim=-1)
        mean_diff = token_mean_diff.mean().item()
        max_diff = (mg_logits - hf_logits).abs().max().item()
        loss_mask = (torch.roll(hf_inputs['labels'], -1) != -100)
        mean_diff_with_loss = token_mean_diff[loss_mask].mean().item()
        max_diff_with_loss = (mg_logits - hf_logits)[loss_mask].abs().max().item()
        print(f'token_mean_diff: {token_mean_diff}')
        print(f'mean_diff: {mean_diff}, max_diff: {max_diff}')
        print(f'mean_diff (with loss): {mean_diff_with_loss}, max_diff (with loss): {max_diff_with_loss} '
              '(Please check that mean_diff (with loss) is less than 0.1).')
        hf_tokens = hf_logits.argmax(-1)
        mg_tokens = mg_logits.argmax(-1)
        print(f'hf_tokens: {hf_tokens[0].tolist()}\nmg_tokens: {mg_tokens[0].tolist()}')
        print(f'token_diff: {(hf_tokens != mg_tokens).sum().item()}')
        print(f'token_diff (with loss): {(hf_tokens[loss_mask] != mg_tokens[loss_mask]).sum().item()}')
