# Copyright (c) ModelScope Contributors. All rights reserved.
# Part of the implementation is borrowed from huggingface/transformers.
import inspect
import math
import os
import torch
import torch.distributed as dist
import torch.nn.functional as F
from contextlib import contextmanager
from modelscope.hub.api import HubApi
from peft import PeftModel
from torch import nn
from torch.nn import CrossEntropyLoss, Module
from transformers import PreTrainedModel
from types import FunctionType, MethodType
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union

from swift.model import ModelMeta
from swift.sequence_parallel import ChunkedCrossEntropyLoss, GatherLoss, sequence_parallel
from swift.utils import deep_getattr, get_dist_setting, get_logger

if TYPE_CHECKING:
    from .arguments import TrainingArguments

logger = get_logger()


def _get_deepspeed_elastic_world_size():
    if dist.is_available() and dist.is_initialized():
        return dist.get_world_size()
    return get_dist_setting()[2]


def _enable_load_universal(ds_config):
    if isinstance(ds_config, dict):
        checkpoint = ds_config.get('checkpoint')
        if not isinstance(checkpoint, dict):
            checkpoint = {}
            ds_config['checkpoint'] = checkpoint
        checkpoint['load_universal'] = True


def enable_deepspeed_load_universal(args: 'TrainingArguments', trainer=None):
    _enable_load_universal(getattr(args, 'deepspeed', None))

    hf_ds_config = getattr(args, 'hf_deepspeed_config', None)
    _enable_load_universal(getattr(hf_ds_config, 'config', None))

    deepspeed_plugin = getattr(args, 'deepspeed_plugin', None)
    if trainer is not None and deepspeed_plugin is None:
        accelerator = getattr(trainer, 'accelerator', None)
        state = getattr(accelerator, 'state', None)
        deepspeed_plugin = getattr(state, 'deepspeed_plugin', None)
    if deepspeed_plugin is not None:
        _enable_load_universal(getattr(deepspeed_plugin, 'deepspeed_config', None))
        plugin_hf_ds_config = getattr(deepspeed_plugin, 'hf_ds_config', None)
        _enable_load_universal(getattr(plugin_hf_ds_config, 'config', None))


def prepare_deepspeed_elastic_config(args: 'TrainingArguments', state=None):
    ds_config = args.deepspeed
    if not ds_config:
        return
    if not isinstance(ds_config, dict):
        logger.warning('DeepSpeed elastic expects args.deepspeed to be a dict, but got '
                       f'{type(ds_config).__name__}. Skip elastic config.')
        return

    from deepspeed.elasticity import compute_elastic_config
    from deepspeed.git_version_info import version as __version__

    enable_deepspeed_load_universal(args)
    elasticity = ds_config.get('elasticity') or {}
    if not elasticity:
        logger.warning_once('DeepSpeed elastic callback is enabled, but no `elasticity` section is found in '
                            'the DeepSpeed config. Only `checkpoint.load_universal` is enabled.')
        return
    if elasticity.get('enabled') is False:
        return

    world_size = _get_deepspeed_elastic_world_size()
    final_batch_size, _, micro_batch_size = compute_elastic_config(
        ds_config=ds_config,
        target_deepspeed_version=__version__,
        world_size=world_size,
    )
    if world_size <= 0 or micro_batch_size <= 0:
        raise ValueError('DeepSpeed elastic config produced invalid batch settings: '
                         f'world_size={world_size}, micro_batch_size={micro_batch_size}.')
    gradient_accu_steps = max(1, final_batch_size // (micro_batch_size * world_size))
    args.per_device_train_batch_size = micro_batch_size
    args.gradient_accumulation_steps = gradient_accu_steps
    if state is not None:
        state.train_batch_size = args.per_device_train_batch_size * max(1, args.n_gpu)
    logger.info_once('DeepSpeed elastic config is enabled. '
                     f'world_size: {world_size}, '
                     f'per_device_train_batch_size: {args.per_device_train_batch_size}, '
                     f'gradient_accumulation_steps: {args.gradient_accumulation_steps}')


def can_return_loss(model: Module) -> bool:
    """Check if a given model can return loss."""
    if isinstance(model, PeftModel):
        signature = inspect.signature(model.model.forward)
    else:
        signature = inspect.signature(model.forward)
    for p in signature.parameters:
        if p == 'return_loss' and signature.parameters[p].default is True:
            return True
    return False


def find_labels(model: Module) -> List[str]:
    """Find the labels used by a given model."""
    model_name = model.__class__.__name__
    if isinstance(model, PeftModel):
        signature = inspect.signature(model.model.forward)
    else:
        signature = inspect.signature(model.forward)
    if 'QuestionAnswering' in model_name:
        return [p for p in signature.parameters if 'label' in p or p in ('start_positions', 'end_positions')]
    else:
        return [p for p in signature.parameters if 'label' in p]


def get_function(method_or_function: Union[MethodType, FunctionType]) -> FunctionType:
    if isinstance(method_or_function, MethodType):
        method_or_function = method_or_function.__func__
    return method_or_function


def is_instance_of_ms_model(model: Module) -> bool:
    """avoid import modelscope: circular dependency problem"""
    for m_cls in model.__class__.__mro__:
        cls_name = m_cls.__name__
        cls_module = m_cls.__module__
        if cls_name == 'Model' and cls_module.startswith('modelscope'):
            return True
    return False


def per_token_loss_func_sp(outputs, labels, enable_dft_loss=False, **kwargs) -> torch.Tensor:
    """Common loss function for sequence parallel training"""
    if hasattr(outputs, 'logits'):
        logits = outputs.logits
    else:
        logits = outputs
    device = logits.device

    batch_size = logits.shape[0]
    logits = logits.view(-1, logits.shape[-1])
    labels = labels.flatten().to(device)
    sploss_parallel_size = int(os.environ.get('CELOSS_PARALLEL_SIZE', '0'))
    if sploss_parallel_size > 0:
        loss = ChunkedCrossEntropyLoss.apply(logits, labels, sploss_parallel_size)
    else:
        loss_fct = CrossEntropyLoss(reduction='none')
        loss = loss_fct(logits, labels)
    if enable_dft_loss:
        with torch.no_grad():
            target_probs = torch.exp(-loss)
        loss *= target_probs
    position_ids = sequence_parallel.real_position_ids
    if position_ids is not None:
        position_ids = sequence_parallel.pad(position_ids, padding_value=-1, position_ids=position_ids)
    loss, labels = GatherLoss.apply(loss.reshape(batch_size, -1), labels.reshape(batch_size, -1), 1, position_ids)
    if position_ids is not None and position_ids.min() == -1:
        _pos_mask = position_ids >= 0
        loss = loss[_pos_mask].contiguous()

    return loss


def per_token_loss_func(outputs, labels, enable_dft_loss: bool = False, **kwargs):
    logits = outputs.logits
    # Upcast to float if we need to compute the loss to avoid potential precision issues
    logits = logits.float()
    labels = torch.roll(labels, shifts=-1, dims=-1).view(-1)

    # Flatten the tokens
    logits = logits.view(-1, logits.shape[-1])
    # Enable model parallelism
    labels = labels.to(logits.device)
    loss = F.cross_entropy(logits, labels, ignore_index=-100, reduction='none')
    if enable_dft_loss:
        with torch.no_grad():
            target_probs = torch.exp(-loss)
        loss *= target_probs
    return loss


def _kwargs_to_args(func, args, kwargs) -> Optional[List[Any]]:
    parameters = inspect.signature(func).parameters
    args = list(args)
    parameters = list(parameters.items())[len(args):]
    for key, param in parameters:
        if key in kwargs:
            args.append(kwargs[key])
        elif param.default != param.empty:
            args.append(param.default)
        else:
            return
    return args


def _add_gradient_checkpointing(module_list):

    requires_grad = None

    def _new_forward(self, *args, **kwargs):
        nonlocal requires_grad
        if requires_grad is None:
            requires_grad = any(p.requires_grad for p in self.parameters())

        new_args = _kwargs_to_args(self.__old_forward, args, kwargs)
        if new_args is not None and self.gradient_checkpointing and self.training:
            if new_args and isinstance(new_args[0], torch.Tensor) and requires_grad and not new_args[0].requires_grad:
                new_args[0].requires_grad_(True)
            layer_ret = self._gradient_checkpointing_func(self.__old_forward, *new_args)
            logger.info_once('Successfully using dynamic gradient checkpointing.')
        else:
            layer_ret = self.__old_forward(*args, **kwargs)
        return layer_ret

    for module in module_list:
        module.gradient_checkpointing = False
        if hasattr(module, '_old_forward'):  # device_map
            __old_forward = module._old_forward
            module._old_forward = MethodType(_new_forward, module)
        else:
            __old_forward = module.forward
            module.forward = MethodType(_new_forward, module)
        module.__old_forward = __old_forward


def find_module_list(model) -> Optional[nn.ModuleList]:
    module_lists = []
    for m in model.modules():
        if hasattr(m, 'gradient_checkpointing') or m.__class__.__name__ == 'CheckpointWrapper':
            return
        if (isinstance(m, (nn.ModuleList, nn.Sequential)) and len(m) >= 10
                and 'mlp' not in m[0].__class__.__name__.lower()):  # fix moe
            module_lists.append(m)
    if module_lists:
        return max(module_lists, key=lambda x: len(x))


def dynamic_gradient_checkpointing(model, including_vit: bool = False) -> None:
    if isinstance(model, PeftModel):
        model = model.model
    model_meta: ModelMeta = getattr(model, 'model_meta', None)
    if model_meta is not None and model_meta.is_multimodal and model_meta.model_arch:
        tower_names = model_meta.model_arch.language_model.copy()
        if including_vit:
            tower_names += model_meta.model_arch.vision_tower
    else:
        tower_names = [None]

    model.supports_gradient_checkpointing = True
    for tower_name in tower_names:
        if tower_name is None:
            model_tower = model
        else:
            model_tower = deep_getattr(model, tower_name)
        if model_tower is None:
            continue
        model_tower.supports_gradient_checkpointing = True
        module_list = find_module_list(model_tower)
        if module_list is None:
            continue
        _add_gradient_checkpointing(module_list)
        logger.info(f'Automatically add gradient_checkpointing to {model_tower.__class__}.')


@contextmanager
def disable_gradient_checkpointing(model: PreTrainedModel, gradient_checkpointing_kwargs: Optional[Dict] = None):
    """
    Temporarily disable gradient checkpointing, restoring the previous state afterward.

    When gradient checkpointing is enabled with use_reentrant=True (default), calling the model inside a
    torch.no_grad() block triggers a harmless PyTorch warning ("None of the inputs have requires_grad=True").
    Temporarily disable checkpointing to avoid this warning during inference.

    Args:
        model (`PreTrainedModel`):
            Model for which to temporarily disable gradient checkpointing.
        gradient_checkpointing_kwargs (`dict` or `None`, *optional*):
            Additional kwargs for gradient checkpointing enabling.
    """
    was_enabled = getattr(model, 'is_gradient_checkpointing', False)
    if was_enabled:
        model.gradient_checkpointing_disable()
    try:
        yield
    finally:
        if was_enabled:
            model.gradient_checkpointing_enable(gradient_checkpointing_kwargs)


def gather_for_unpadded_tensors(input_data, use_gather_object=False):
    from accelerate.utils import gather_object
    if getattr(sequence_parallel, 'dp_group', None) is not None:
        input_data = sequence_parallel._gather_object_dp(input_data)
    else:
        input_data = gather_object(input_data)
    output = []
    for _data in input_data:
        if len(_data.shape) == 0:
            _data = _data.unsqueeze(0)
        _data = _data.cpu()
        output.append(_data)
    if len(output[0].shape) == 1 and output[0].shape[0] > 1:
        data = torch.stack(output, dim=0)
    else:
        data = torch.concat(output, dim=0)
    return data


def calculate_max_steps(args: 'TrainingArguments', dataset) -> int:
    if args.max_steps and args.max_steps > 0:
        max_steps = args.max_steps
    else:
        len_dataset = len(dataset)
        _, _, world_size, _ = get_dist_setting()
        total_train_batch_size = args.per_device_train_batch_size * args.gradient_accumulation_steps * world_size
        num_update_steps_per_epoch = len_dataset // total_train_batch_size
        num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1)
        max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch)
    return max_steps


def extract_version(name: str) -> Optional[int]:
    if not name.startswith('v'):
        return None
    try:
        num = name[1:].split('-', 1)[0]
        return int(num)
    except ValueError:
        return None


def get_previous_version_from_path(current_path: str) -> Optional[str]:
    from pathlib import Path
    current = Path(current_path)
    parent = current.parent
    current_name = current.name

    candidates = [d for d in parent.iterdir() if d.is_dir()]

    valid = [(d.name, extract_version(d.name)) for d in candidates]
    valid = [(name, ver) for name, ver in valid if ver is not None]

    valid.sort(key=lambda x: x[1])
    names = [name for name, _ in valid]

    if current_name not in names:
        return None

    idx = names.index(current_name)
    if idx == 0:
        return None

    prev_name = names[idx - 1]
    return str(parent / prev_name)


def get_resume_dir(output_dir):
    return get_previous_version_from_path(output_dir)


def replace_index_file(output_dir: str):
    import json
    import os
    from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, WEIGHTS_INDEX_NAME
    index_file = os.path.join(output_dir, WEIGHTS_INDEX_NAME)

    if not os.path.exists(index_file):
        return
    with open(index_file, 'r', encoding='utf-8') as f:
        bin_data = json.load(f)
    if 'weight_map' not in bin_data:
        return
    bin_data['weight_map'] = {
        k: v.replace('pytorch_model', 'model').replace('.bin', '.safetensors')
        for k, v in bin_data['weight_map'].items()
    }
    safe_path = os.path.join(output_dir, SAFE_WEIGHTS_INDEX_NAME)
    with open(safe_path, 'w', encoding='utf-8') as f:
        json.dump(bin_data, f, indent=2)
    from contextlib import suppress
    with suppress(FileNotFoundError):
        os.remove(os.path.join(output_dir, WEIGHTS_INDEX_NAME))


@contextmanager
def patch_modelscope_hub_timeout():
    __init__ = HubApi.__init__

    def __new_init__(self, *args, **kwargs):
        timeout = kwargs.get('timeout')
        if timeout is not None and timeout > 5:
            kwargs['timeout'] = 5
        __init__(self, *args, **kwargs)

    HubApi.__init__ = __new_init__

    try:
        yield
    finally:
        HubApi.__init__ = __init__