import torch
from tqdm.auto import tqdm
from safetensors.torch import safe_open
from typing import List, Tuple, Generator
from vllm.model_executor.model_loader.loader import DefaultModelLoader
from vllm.model_executor.model_loader.weight_utils import (np_cache_weights_iterator,
                                                           _BAR_FORMAT)
from vllm.config import LoadFormat
from vllm.platforms import current_platform
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm_mlu._mlu_utils import get_device_major_capability


CAST_BFLOAT16_TO_FLOAT16_ENABLE = (get_device_major_capability() == 3)


def vllm__model_executor__model_loader__weight_utils__safetensors_weights_iterator(
    hf_weights_files: List[str]
) -> Generator[Tuple[str, torch.Tensor], None, None]:
    """Iterate over the weights in the model safetensor files."""
    enable_tqdm = not torch.distributed.is_initialized(
    ) or torch.distributed.get_rank() == 0
    for st_file in tqdm(
            hf_weights_files,
            desc="Loading safetensors checkpoint shards",
            disable=not enable_tqdm,
            bar_format=_BAR_FORMAT,
    ):
        with safe_open(st_file, framework="pt") as f:
            for name in f.keys():  # noqa: SIM118
                param = f.get_tensor(name)
                '''
                =============================
                Modify by vllm_mlu
                =============================
                @brief: cast bfloat16 to float16 for MLU3xx
                '''
                if CAST_BFLOAT16_TO_FLOAT16_ENABLE and param.dtype == torch.bfloat16:
                    param = param.to(torch.float16)
                '''
                ==================
                End of MLU Hijack
                ==================
                '''
                yield name, param


def vllm__model_executor__model_loader__weight_utils__pt_weights_iterator(
    hf_weights_files: List[str]
) -> Generator[Tuple[str, torch.Tensor], None, None]:
    """Iterate over the weights in the model bin/pt files."""
    enable_tqdm = not torch.distributed.is_initialized(
    ) or torch.distributed.get_rank() == 0
    for bin_file in tqdm(
            hf_weights_files,
            desc="Loading pt checkpoint shards",
            disable=not enable_tqdm,
            bar_format=_BAR_FORMAT,
    ):
        state = torch.load(bin_file, map_location="cpu")
        for name, param in state.items():
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: cast bfloat16 to float16 for MLU3xx
            '''
            if CAST_BFLOAT16_TO_FLOAT16_ENABLE and param.dtype == torch.bfloat16:
                param = param.to(torch.float16)
            '''
            ==================
            End of MLU Hijack
            ==================
            '''
            yield name, param
        del state
        torch.mlu.empty_cache()


def vllm__model_executor__model_loader__loader__DefaultModelLoader___get_weights_iterator(
        self, source: "DefaultModelLoader.Source"
) -> Generator[Tuple[str, torch.Tensor], None, None]:
    """Get an iterator for the model weights based on the load format."""
    hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
        source.model_or_path, source.revision, source.fall_back_to_pt)
    if self.load_config.load_format == LoadFormat.NPCACHE:
        # Currently np_cache only support *.bin checkpoints
        assert use_safetensors is False
        weights_iterator = np_cache_weights_iterator(
            source.model_or_path, self.load_config.download_dir, hf_folder,
            hf_weights_files)
    elif use_safetensors:
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: cast bfloat16 to float16 for MLU3xx
        '''
        weights_iterator = vllm__model_executor__model_loader__weight_utils__safetensors_weights_iterator(hf_weights_files)
        '''
        ==================
        End of MLU Hijack
        ==================
        '''
    else:
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: cast bfloat16 to float16 for MLU3xx
        '''
        weights_iterator = vllm__model_executor__model_loader__weight_utils__pt_weights_iterator(hf_weights_files)
        '''
        ==================
        End of MLU Hijack
        ==================
        '''

    if current_platform.is_tpu():
        # In PyTorch XLA, we should call `xm.mark_step` frequently so that
        # not too many ops are accumulated in the XLA program.
        import torch_xla.core.xla_model as xm

        def _xla_weights_iterator(iterator: Generator):
            for weights in iterator:
                yield weights
                xm.mark_step()

        weights_iterator = _xla_weights_iterator(weights_iterator)

    # Apply the prefix.
    return ((source.prefix + name, tensor)
            for (name, tensor) in weights_iterator)


MluHijackObject.apply_hijack(DefaultModelLoader,
                             DefaultModelLoader._get_weights_iterator,
                             vllm__model_executor__model_loader__loader__DefaultModelLoader___get_weights_iterator)
