# Copyright (c) ModelScope Contributors. All rights reserved.
import ast
import math
import os
import torch
from dataclasses import dataclass, field
from transformers.utils import is_torch_mps_available
from typing import Any, Dict, List, Literal, Optional, Union

from swift.model import MODEL_MAPPING, get_model_info_meta, get_model_name
from swift.utils import HfConfigFactory, get_dist_setting, get_logger, json_parse_to_dict

logger = get_logger()


@dataclass
class ModelArguments:
    """A dataclass that holds various arguments related to model configuration and usage.

    Args:
        model (Optional[str]): The model ID from the Hub or a local path to the model. Defaults to None.
        model_type (Optional[str]): The model type. In ms-swift, a 'model_type' groups models with the same
            architecture, loading process, and template. Defaults to None, which enables auto-selection based on
            the suffix of `--model` and the 'architectures' attribute in `config.json`. The `model_type` for a
            corresponding model can be found in the list of supported models. Note: The concept of `model_type`
            in ms-swift differs from the `model_type` in `config.json`. Custom models usually require registering
            their own `model_type` and `template`.
        model_revision (Optional[str]): The revision of the model. Defaults to None.
        task_type (str): The task type. Can be 'causal_lm', 'seq_cls', 'embedding', 'reranker', or
            'generative_reranker'. If set to 'seq_cls', you usually need to specify `--num_labels` and
            `--problem_type`. Defaults to 'causal_lm'.
        torch_dtype (Optional[str]): The data type of the model weights. Supports 'float16', 'bfloat16', 'float32'.
            Defaults to None, in which case it's read from the 'config.json' file.
        attn_impl (Optional[str]): The attention implementation to use. Options include 'sdpa', 'eager', 'flash_attn',
            'flash_attention_2', 'flash_attention_3', etc. Defaults to None, which means it will be read from
            'config.json'. Note: Support for these implementations depends on the model's transformers implementation.
            If set to 'flash_attn' (for backward compatibility), 'flash_attention_2' will be used.
        experts_impl (Optional[str]): Expert implementation type, options are 'grouped_mm', 'batched_mm', 'eager'.
            Defaults to None. This feature requires "transformers>=5.0.0".
        new_special_tokens (List[str]): Additional special tokens to be added to the tokenizer. Can also be a path to
            a `.txt` file, where each line is a special token. Defaults to an empty list `[]`.
        num_labels (Optional[int]): The number of labels for classification tasks (when `--task_type` is 'seq_cls').
            Required for such tasks. Defaults to None.
        problem_type (Optional[str]): The problem type for classification tasks (`--task_type` 'seq_cls'). Options are
            'regression', 'single_label_classification', 'multi_label_classification'. Defaults to None, but is
            automatically set to 'regression' if the model is a reward_model or `num_labels` is 1, and
            'single_label_classification' otherwise.
        rope_scaling (Optional[str]): The RoPE scaling type. You can pass a string like 'linear', 'dynamic', or
            'yarn', and ms-swift will automatically set the corresponding `rope_scaling` and override the
            'config.json' value. Alternatively, you can pass a JSON string (e.g., '{"factor":2.0, "type":"yarn"}'),
            which will directly override the `rope_scaling` in 'config.json'. Defaults to None.
        device_map (Optional[str]): The device map configuration for the model, e.g., 'auto', 'cpu', a JSON string,
            or a path to a JSON file. This argument is passed directly to the `from_pretrained` method of transformers.
            Defaults to None, and will be set automatically based on the device and distributed training settings.
        max_memory (Optional[str]): The maximum memory allocation for each device when `device_map` is 'auto' or
            'sequential'. Example: '{0: "20GB", 1: "20GB"}'. This argument is passed directly to the `from_pretrained`
            method of transformers. Defaults to None.
        max_model_len (Optional[int]): The maximum model length. This is used to calculate the RoPE scaling factor
            when `rope_scaling` is specified as a string. If not None, it overrides the `max_position_embeddings`
            value in 'config.json'. Defaults to None.
        local_repo_path (Optional[str]): Path to a local repository for models that require a GitHub repo during
            loading (e.g., deepseek-vl2). This avoids network issues during `git clone`. Defaults to None.
        init_strategy (Optional[str]): The strategy to initialize all uninitialized parameters when loading a model
            (especially for custom architectures). Options include 'zero', 'uniform', 'normal', 'xavier_uniform',
            'xavier_normal', 'kaiming_uniform', 'kaiming_normal', 'orthogonal'. Defaults to None.
    """
    model: Optional[str] = None  # model id or model path
    model_type: Optional[str] = field(
        default=None, metadata={'help': f'model_type choices: {list(MODEL_MAPPING.keys())}'})
    model_revision: Optional[str] = None
    task_type: Literal['causal_lm', 'seq_cls', 'embedding', 'reranker', 'generative_reranker'] = None

    torch_dtype: Literal['bfloat16', 'float16', 'float32', None] = None
    # flash_attn: It will automatically convert names based on the model.
    # None: It will be automatically selected between sdpa and eager.
    # 'flash_attn', 'sdpa', 'eager', 'flex_attention', 'flash_attention_2', 'flash_attention_3'
    attn_impl: Optional[str] = None
    experts_impl: Optional[str] = None
    new_special_tokens: List[str] = field(default_factory=list)

    num_labels: Optional[int] = None
    problem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] = None
    rope_scaling: Optional[str] = None
    device_map: Optional[Union[dict, str]] = None
    max_memory: Optional[Union[dict, str]] = None
    max_model_len: Optional[int] = None
    # When some model code needs to be downloaded from GitHub,
    # this parameter specifies the path to the locally downloaded repository.
    local_repo_path: Optional[str] = None
    init_strategy: Literal['zero', 'uniform', 'normal', 'xavier_uniform', 'xavier_normal', 'kaiming_uniform',
                           'kaiming_normal', 'orthogonal'] = None

    def _init_device_map(self):
        """Prepare device map args"""
        if self.device_map:
            self.device_map: Union[str, Dict[str, Any], None] = json_parse_to_dict(self.device_map, strict=False)
        # compat mp&ddp
        _, local_rank, _, local_world_size = get_dist_setting()
        if local_world_size > 1 and isinstance(self.device_map, dict) and local_rank > 0:
            for k, v in self.device_map.items():
                if isinstance(v, int):
                    self.device_map[k] += local_rank

    def _init_max_memory(self):
        if isinstance(self.max_memory, str):
            try:
                self.max_memory = ast.literal_eval(self.max_memory)
            except Exception:
                pass
        self.max_memory = json_parse_to_dict(self.max_memory)
        # compat mp&ddp
        _, local_rank, _, local_world_size = get_dist_setting()
        if local_world_size > 1 and isinstance(self.max_memory, dict) and local_rank > 0:
            for k in list(self.max_memory.keys()):
                if isinstance(k, int):
                    self.max_memory[k + local_rank] = self.max_memory.pop(k)

    def _init_torch_dtype(self) -> None:
        """"If torch_dtype is None, find a proper dtype by the config.json/GPU"""
        from ..sft_args import SftArguments

        self.torch_dtype: Optional[torch.dtype] = HfConfigFactory.to_torch_dtype(self.torch_dtype)
        self.torch_dtype: torch.dtype = self._init_model_info()
        # Mixed Precision Training
        if isinstance(self, SftArguments):
            self._init_mixed_precision()

    def _init_mixed_precision(self):
        if is_torch_mps_available():
            fp16, bf16 = False, False
        elif self.torch_dtype in {torch.float16, torch.float32}:
            fp16, bf16 = True, False
        elif self.torch_dtype == torch.bfloat16:
            fp16, bf16 = False, True
        else:
            raise ValueError(f'args.torch_dtype: {self.torch_dtype}')
        if self.fp16 is None:
            self.fp16 = fp16
        if self.bf16 is None:
            self.bf16 = bf16

    def _init_rope_scaling(self):
        if self.rope_scaling:
            rope_scaling: dict = json_parse_to_dict(self.rope_scaling, strict=False)
            if isinstance(rope_scaling, str):
                assert rope_scaling in ['linear', 'dynamic', 'yarn']
                rope_scaling = {'type': rope_scaling}
        else:
            rope_scaling = self.model_info.rope_scaling
            # reset the factor
            rope_scaling.pop('factor', None)

        rope_type = rope_scaling.get('rope_type', rope_scaling.get('type', 'default'))
        if 'factor' not in rope_scaling and self.max_model_len is None and rope_type == 'default':
            # fix megatron qwen2_5_vl
            self.rope_scaling = rope_scaling
            logger.info(f'Setting args.rope_scaling: {rope_scaling}')
            return

        # get origin_max_model_len
        origin_max_model_len = None
        if rope_scaling and rope_scaling.get('original_max_position_embeddings') is not None:
            origin_max_model_len = rope_scaling['original_max_position_embeddings']
        elif self.model_info.rope_scaling:
            if self.model_info.rope_scaling.get('original_max_position_embeddings') is not None:
                origin_max_model_len = self.model_info.rope_scaling['original_max_position_embeddings']
            elif self.model_info.rope_scaling.get('factor') is not None:
                origin_max_model_len = self.model_info.max_model_len // self.model_info.rope_scaling['factor']
        if origin_max_model_len is None:
            origin_max_model_len = self.model_info.max_model_len
        assert origin_max_model_len is not None, '`origin_max_model_len` from model config is not set'
        rope_scaling['original_max_position_embeddings'] = origin_max_model_len

        if 'factor' not in rope_scaling:
            assert self.max_model_len is not None, (
                'max_model_len must be set if rope_scaling does not contain a "factor"')
            rope_scaling['factor'] = max(float(math.ceil(self.max_model_len / origin_max_model_len)), 1.0)
        rope_model_len = int(origin_max_model_len * rope_scaling['factor'])
        if self.max_model_len is None:
            self.max_model_len = rope_model_len
        elif self.max_model_len > rope_model_len:
            logger.warning(f'rope config ({rope_model_len} = {rope_scaling["factor"]} * '
                           f'{origin_max_model_len}) should be bigger than max_model_len '
                           f'from command line ({self.max_model_len})')
        self.rope_scaling = rope_scaling
        logger.info(f'Setting args.rope_scaling: {rope_scaling}')
        logger.info(f'Setting args.max_model_len: {self.max_model_len}')

    def _init_model_info(self) -> torch.dtype:
        model_kwargs = self.get_model_kwargs()
        if self.tuner_backend == 'unsloth':
            model_kwargs['download_model'] = True
        self.model_info, self.model_meta = get_model_info_meta(**model_kwargs)
        self.task_type = self.model_info.task_type
        self.num_labels = self.model_info.num_labels

        self.model_dir = self.model_info.model_dir
        self.model_type = self.model_info.model_type
        if self.rope_scaling or self.model_info.rope_scaling and self.max_model_len is not None:
            self._init_rope_scaling()
        return self.model_info.torch_dtype

    def _init_new_special_tokens(self):
        if isinstance(self.new_special_tokens, str):
            self.new_special_tokens = [self.new_special_tokens]
        new_special_tokens = []
        for token in self.new_special_tokens:
            if token.endswith('.txt'):
                assert os.path.isfile(token), f'special_tokens_path: {token}'
                with open(token, 'r') as f:
                    text = f.read()
                new_special_tokens += text.split()
            else:
                new_special_tokens.append(token)
        self.new_special_tokens = new_special_tokens

    def __post_init__(self):
        if self.model is None:
            raise ValueError(f'Please set --model <model_id_or_path>`, model: {self.model}')
        self._init_new_special_tokens()
        self.model_suffix = get_model_name(self.model)
        self._init_device_map()
        self._init_max_memory()
        self._init_torch_dtype()

    def get_model_kwargs(self):
        return {
            'model_id_or_path': self.model,
            'torch_dtype': self.torch_dtype,
            'model_type': self.model_type,
            'revision': self.model_revision,
            'use_hf': self.use_hf,
            'hub_token': self.hub_token,
            'local_repo_path': self.local_repo_path,
            'device_map': self.device_map,
            'max_memory': self.max_memory,
            'quantization_config': self.get_quantization_config(),
            'attn_impl': self.attn_impl,
            'new_special_tokens': self.new_special_tokens,
            'rope_scaling': self.rope_scaling,
            'max_model_len': self.max_model_len,
            'task_type': self.task_type,
            'num_labels': self.num_labels,
            'problem_type': self.problem_type,
            'init_strategy': self.init_strategy,
        }
