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

import math
import os
import shutil
import torch
from transformers.utils import strtobool

from swift.arguments import ExportArguments
from swift.pipelines import prepare_model_template
from swift.utils import get_logger, get_n_params_grads, is_master
from .arguments import MegatronArguments
from .model import get_mcore_model
from .utils import (load_mcore_checkpoint, patch_torch_dist_shard, prepare_mcore_model, save_mcore_checkpoint,
                    test_convert_precision)

logger = get_logger()

convert_kwargs = {
    'use_cpu_initialization': True,
    'no_save_optim': True,
    'no_save_rng': True,
    'no_load_optim': True,
    'no_load_rng': True,
    'finetune': True,
    'attention_backend': 'unfused',
    'padding_free': False,
    'recompute_granularity': 'none',  # deepseek-v4
}


def convert_hf2mcore(args: ExportArguments) -> None:
    args.experts_impl = 'eager'  # Compatible with transformers 5.4.0
    hf_model, template = prepare_model_template(args, patch_offload=not args.test_convert_precision)
    processor = template.processor
    if args.thread_count is None:
        checkpoint_size = sum(get_n_params_grads(hf_model)[0]) * torch.finfo(args.torch_dtype).bits // 8e9
        args.thread_count = max(math.ceil(checkpoint_size / 10), 2)  # 10GB
    patch_torch_dist_shard(args.thread_count)

    hf_config = processor.model_info.config
    current_convert_kwargs = convert_kwargs.copy()
    if args.model_info.is_moe_model:
        current_convert_kwargs['moe_grouped_gemm'] = True
    megatron_args = MegatronArguments(
        model=args.model,
        model_type=args.model_type,
        **current_convert_kwargs,
        output_dir=args.output_dir,
        torch_dtype=args.torch_dtype)

    mg_model = get_mcore_model(megatron_args, hf_config)[0]
    logger.info('Megatron model created successfully.')
    bridge = mg_model.config.bridge
    bridge.load_weights([mg_model], args.model_info.model_dir)
    logger.info('Successfully transferred HF model weights to MG model.')
    _test_convert_precision = strtobool(os.getenv('SWIFT_TEST_CONVERT_PRECISION', '0'))
    if not _test_convert_precision:
        args.save_args()
        logger.info('Saving the model...')
        save_mcore_checkpoint(megatron_args, [mg_model])
    # Place it at the end to avoid test_convert_precision affecting precision.
    if args.test_convert_precision:
        test_convert_precision(megatron_args, hf_model, mg_model, template, test_convert_dtype=args.test_convert_dtype)


def convert_mcore2hf(args: ExportArguments) -> None:
    args.experts_impl = 'eager'
    _, template = prepare_model_template(args, load_model=False)
    processor = template.processor

    hf_config = processor.model_info.config
    current_convert_kwargs = convert_kwargs.copy()
    if args.model_info.is_moe_model:
        current_convert_kwargs['moe_grouped_gemm'] = True
    extra_config = MegatronArguments.load_args_config(args.mcore_adapter or args.mcore_model)
    extra_config['mcore_adapter'] = args.mcore_adapter
    if args.mcore_model is not None:
        extra_config['mcore_model'] = args.mcore_model
    current_convert_kwargs.update(extra_config)
    megatron_args = MegatronArguments(
        model=args.model,
        model_type=args.model_type,
        **current_convert_kwargs,
        output_dir=args.output_dir if args.to_mcore else None,
        torch_dtype=args.torch_dtype)

    mg_model = get_mcore_model(megatron_args, hf_config)[0]
    if megatron_args.mcore_model is None:
        raise ValueError('Please specify `--mcore_model`.')
    load_mcore_checkpoint(megatron_args, [mg_model], load_arg='mcore_model')
    if megatron_args.mcore_adapter is not None:
        peft_model = prepare_mcore_model(megatron_args, mg_model)
        load_mcore_checkpoint(megatron_args, [mg_model], load_arg='mcore_adapter')
        logger.info('Merge LoRA...')
        mg_model = peft_model.merge_and_unload()
    logger.info('Megatron model created successfully.')
    if args.to_hf:
        bridge = mg_model.config.bridge
        logger.info('Converting weights and saving the model...')
        bridge.save_weights([mg_model], args.output_dir, args=megatron_args, processor=processor)
        if is_master():
            if args.ckpt_dir:
                args_path = os.path.join(args.ckpt_dir, 'args.json')
                shutil.copy(args_path, os.path.join(args.output_dir, 'args.json'))
            else:
                args.save_args(args.output_dir)
        if args.test_convert_precision:
            hf_model, template = prepare_model_template(args, model=args.output_dir)
            test_convert_precision(
                megatron_args, hf_model, mg_model, template, test_convert_dtype=args.test_convert_dtype)
    elif args.to_mcore:
        if args.thread_count is None:
            checkpoint_size = sum(get_n_params_grads(mg_model)[0]) * torch.finfo(args.torch_dtype).bits // 8e9
            args.thread_count = max(math.ceil(checkpoint_size / 10), 2)  # 10GB
        patch_torch_dist_shard(args.thread_count)

        args.save_args()
        logger.info('Saving the model...')
        save_mcore_checkpoint(megatron_args, [mg_model])