import torch
import gc

import functools
from collections import defaultdict
from typing import Tuple

from vllm_mlu.model_executor.layers.feed_forward import FeedForward
from vllm.model_executor.layers.vocab_parallel_embedding import (VocabParallelEmbedding,
                                                                 ParallelLMHead)
from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
    MergedColumnParallelLinear,
    QKVParallelLinear,
    RowParallelLinear)
from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp

from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.model_executor import set_random_seed
from vllm.worker.mlu_worker import MLUWorker
from vllm.logger import init_logger

logger = init_logger(__name__)

def default_act_range_value():
    return {
        "x": None,
        "split": None,
        "is_linear": False,
        "is_qkv": False,
        "q_proj_size": 0,
        "num_kv_head_replicas": 1,
        "is_merge": False,
        "input_id": []
    }

class MLUWorker_V2(MLUWorker):

    @torch.inference_mode()
    def determine_num_available_blocks(self) -> Tuple[int, int]:
        """Profiles the peak memory usage of the model to determine how many
        KV blocks may be allocated without OOMs.

        The engine will first conduct a profiling of the existing memory usage.
        Then, it calculate the maximum possible number of GPU and CPU blocks
        that can be allocated with the remaining free memory.

        .. tip::
            You may limit the usage of GPU memory
            by adjusting the `gpu_memory_utilization` parameter.
        """
        # Profile the memory usage of the model and get the maximum number of
        # cache blocks that can be allocated with the remaining free memory.
        torch.mlu.empty_cache()
        torch.mlu.reset_peak_memory_stats()

        free_memory_pre_profile, total_gpu_memory = torch.mlu.mem_get_info()

        # Execute a forward pass with dummy inputs to profile the memory usage
        # of the model.
        self.model_runner.profile_run()
        torch.mlu.synchronize()

        self._assert_memory_footprint_increased_during_profiling()

        # Get the peak memory allocation recorded by torch
        peak_memory = torch.mlu.memory_stats()["allocated_bytes.all.peak"]

        # Check for any memory left around that may have been allocated on the
        # gpu outside of `torch`. NCCL operations, for example, can use a few
        # GB during a forward pass
        torch.mlu.empty_cache()
        torch_allocated_bytes = torch.mlu.memory_stats(
        )["allocated_bytes.all.current"]
        total_allocated_bytes = torch.mlu.mem_get_info(
        )[1] - torch.mlu.mem_get_info()[0]
        non_torch_allocations = total_allocated_bytes - torch_allocated_bytes
        if non_torch_allocations > 0:
            peak_memory += non_torch_allocations

        available_kv_cache_memory = (
            total_gpu_memory * self.cache_config.gpu_memory_utilization -
            peak_memory)

        # Calculate the number of blocks that can be allocated with the
        # profiled peak memory.
        cache_block_size = self.get_cache_block_size_bytes()
        if cache_block_size == 0:
            num_gpu_blocks = 0
            num_cpu_blocks = 0
        else:
            num_gpu_blocks = int(available_kv_cache_memory // cache_block_size)
            num_cpu_blocks = int(self.cache_config.swap_space_bytes //
                                    cache_block_size)
        num_gpu_blocks = max(num_gpu_blocks, 0)
        num_cpu_blocks = max(num_cpu_blocks, 0)

        logger.info(
            "Memory profiling results: total_gpu_memory=%.2fGiB"
            " initial_memory_usage=%.2fGiB peak_torch_memory=%.2fGiB"
            " memory_usage_post_profile=%.2fGiB"
            " non_torch_memory=%.2fGiB kv_cache_size=%.2fGiB"
            " gpu_memory_utilization=%.2f", total_gpu_memory / (1024**3),
            (total_gpu_memory - free_memory_pre_profile) / (1024**3),
            (peak_memory - non_torch_allocations) / (1024**3),
            total_allocated_bytes / (1024**3),
            non_torch_allocations / (1024**3),
            available_kv_cache_memory / (1024**3),
            self.cache_config.gpu_memory_utilization)

        # Final cleanup
        if self.model_runner.lora_manager:
            self.model_runner.remove_all_loras()
        gc.collect()

        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: record init memory usage
        ''' 
        # Record memory usage
        self.peak_memory = peak_memory
        self.block_memory = available_kv_cache_memory
        self.num_gpu_blocks = num_gpu_blocks
        self.num_cpu_blocks = num_cpu_blocks
        '''
        ==================
        End of MLU Hijack
        ==================
        '''

        return num_gpu_blocks, num_cpu_blocks

    def _warm_up_model(self) -> None:
        if not self.model_config.enforce_eager:
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: support context mlugraph
            ''' 
            if self.model_config.use_context_mlugraph():
                # Capture MLUGraph both prefill and decode
                self.model_runner.capture_model_with_context(self.gpu_cache,
                                                            self.cache_config.num_gpu_blocks)
            else:
                # Capture MLUGraph only decode
                self.model_runner.capture_model(self.gpu_cache,
                                                self.cache_config.num_gpu_blocks)
            '''
            ==================
            End of MLU Hijack
            ==================
            '''
        # Reset the seed to ensure that the random state is not affected by
        # the model initialization and profiling.
        set_random_seed(self.model_config.seed)

    def get_latency(self):
        start, end = self.model_runner.time_markers
        return start.elapsed_time(end)

    def get_memory_usage(self):
        return (self.peak_memory, self.block_memory,
                self.num_gpu_blocks, self.num_cpu_blocks)

    def recapture_model(
        self,
        context_batch_size_to_capture,
        context_seq_len_to_capture
    ) -> None:
        # Reset history capture context
        self.model_runner.reset_capture_context()

        # Re-capture context and decoder mlugraph
        self.model_runner.model_config.context_batch_size_to_capture = context_batch_size_to_capture
        self.model_runner.model_config.context_seq_len_to_capture = context_seq_len_to_capture
        self._warm_up_model()

    def stat_tensor(self, name, tensor, act_range, key, dim):
        logger.debug(f"name:{name}, key:{key}, dim:{dim}, tensor.shape:{tensor.shape}")
        hidden_dim = tensor.shape[-1]
        # TODO
        # For torch.max has bug which generates nan/inf, so load the tensor to cpu to do torch.max.
        # And need to convert to mlu after the bug is fixed.
        # The pytorch jira: http://jira.cambricon.com/browse/PYTORCH-12199
        tensor = tensor.view(-1, hidden_dim).abs()
        comming_max = torch.max(tensor, dim=dim)[0].float()

        if act_range[name][key] is None:
            act_range[name][key] = comming_max
        else:
            act_range[name][key] = torch.max(act_range[name][key], comming_max)

    def stat_input_hook(self, m, x, y, name, act_range, is_linear, is_save_input_id):
        if isinstance(x, tuple):
            x = x[0]
        if isinstance(y, tuple):
            y = y[0]
        logger.debug(f"name:{name}, x.shape:{x.shape}, y.shape:{y.shape}, m.weight.shape:{m.weight.shape}")
        if is_linear:
            self.stat_tensor(name, x, act_range, "x", 0)
            if act_range[name]["is_qkv"] and is_save_input_id and ".0." in name:
                x_cpu = x.clone().to("cpu")
                act_range[name]["input_id"].append(x_cpu)

    def setup_smooth_hook(self, is_save_input_id: bool = False):
        model = self.model_runner.model
        self.act_range = defaultdict(default_act_range_value)
        self.hooks = []
        linear_class_list = (ColumnParallelLinear, MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear)
        other_class_list = (VocabParallelEmbedding, ParallelLMHead)
        class_list = linear_class_list + other_class_list
        row_class_list = (RowParallelLinear)

        for name, m in model.named_modules():
            if isinstance(m, FeedForward):
                m.use_bt_ffn = False
            if isinstance(m, SparseMoeMlp):
                m.is_use_fused_moe = False

            if isinstance(m, class_list):
                is_linear = True if isinstance(m, linear_class_list) else False
                split_type = "row" if isinstance(m, row_class_list) else "col"
                self.act_range[name]["split"] = split_type
                self.act_range[name]["is_linear"] = is_linear
                if isinstance(m, QKVParallelLinear):
                    self.act_range[name]["is_qkv"] = True
                    self.act_range[name]["q_proj_size"] = m.num_heads * m.head_size
                    self.act_range[name]["num_kv_head_replicas"] = m.num_kv_head_replicas
                self.act_range[name]["is_merge"] = isinstance(m, MergedColumnParallelLinear)

                logger.info(f"rank:{self.rank}, add hook to {name}, is_linear:{is_linear}, split_type:{split_type}")
                self.hooks.append(m.register_forward_hook(functools.partial(self.stat_input_hook,
                                                                            name=name, act_range=self.act_range,
                                                                            is_linear=is_linear,
                                                                            is_save_input_id=is_save_input_id)))

    def remove_hooks(self):
        for h in self.hooks:
            h.remove()

    def get_act_range(self):
        act_range = defaultdict(default_act_range_value)
        for layer_name, layer_range in self.act_range.items():
            for tensor_key, tensor_value in layer_range.items():
                if isinstance(tensor_value, torch.Tensor):
                    act_range[layer_name][tensor_key] = tensor_value.to("cpu")
                elif tensor_key == "input_id" and isinstance(tensor_value, list):
                    input_id_len = len(tensor_value)
                    for i in range(input_id_len):
                        if isinstance(tensor_value[i], torch.Tensor):
                            act_range[layer_name][tensor_key].append(tensor_value[i].to("cpu"))
                        else:
                            act_range[layer_name][tensor_key].append(tensor_value[i])
                else:
                    act_range[layer_name][tensor_key] = tensor_value

        return act_range

    @torch.no_grad()
    def get_named_parameters(self):
        name_parameters = {}
        for name, param in self.model_runner.model.named_parameters():
            name_parameters[name] = param.to("cpu")

        return name_parameters


MluHijackObject.apply_hijack(MLUWorker,
                             MLUWorker.determine_num_available_blocks,
                             MLUWorker_V2.determine_num_available_blocks)
MluHijackObject.apply_hijack(MLUWorker,
                             MLUWorker._warm_up_model,
                             MLUWorker_V2._warm_up_model)
MluHijackObject.apply_hijack(MLUWorker,
                             "get_latency",
                             MLUWorker_V2.get_latency)
MluHijackObject.apply_hijack(MLUWorker,
                             "get_memory_usage",
                             MLUWorker_V2.get_memory_usage)
MluHijackObject.apply_hijack(MLUWorker,
                             "recapture_model",
                             MLUWorker_V2.recapture_model)
MluHijackObject.apply_hijack(MLUWorker,
                             "stat_tensor",
                             MLUWorker_V2.stat_tensor)
MluHijackObject.apply_hijack(MLUWorker,
                             "stat_input_hook",
                             MLUWorker_V2.stat_input_hook)
MluHijackObject.apply_hijack(MLUWorker,
                             "setup_smooth_hook",
                             MLUWorker_V2.setup_smooth_hook)
MluHijackObject.apply_hijack(MLUWorker,
                             "remove_hooks",
                             MLUWorker_V2.remove_hooks)
MluHijackObject.apply_hijack(MLUWorker,
                             "get_act_range",
                             MLUWorker_V2.get_act_range)
MluHijackObject.apply_hijack(MLUWorker,
                             "get_named_parameters",
                             MLUWorker_V2.get_named_parameters)