from typing import Optional, List, Any, Tuple
import torch
from torch.nn.parameter import Parameter, UninitializedParameter

from vllm.distributed import (divide, split_tensor_along_last_dim)
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm.model_executor.utils import set_weight_attrs
from vllm.model_executor.parameter import (BasevLLMParameter,
                                           PerTensorScaleParameter,
                                           RowvLLMParameter)

from vllm.logger import init_logger
from vllm_mlu._mlu_utils import *
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.model_executor.layers.linear import (WEIGHT_LOADER_V2_SUPPORTED, LinearBase, ColumnParallelLinear,
                                               MergedColumnParallelLinear, RowParallelLinear, adjust_marlin_shard,
                                               adjust_scalar_to_fused_array)
from vllm import _mlu_ops as mlu_ops
from ....mlu_hijack.distributed.parallel_state import (get_parallel_rank_with_group, get_parallel_world_size_with_group,
                                                       get_tp_group)
from ....mlu_hijack.distributed.communication_op import (tensor_model_parallel_all_reduce,
                                                         tensor_model_parallel_all_gather)

vllm__model_executor__layers__linear__LinearBase____init__org = LinearBase.__init__

logger = init_logger(__name__)


def vllm__model_executor__layers__linear__LinearBase____init__(
        self,
        input_size: int,
        output_size: int,
        skip_bias_add: bool = False,
        params_dtype: Optional[torch.dtype] = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        tp_group: Any = None,
    ):
    vllm__model_executor__layers__linear__LinearBase____init__org(self=self,
                                                                  input_size=input_size,
                                                                  output_size=output_size,
                                                                  skip_bias_add=skip_bias_add,
                                                                  params_dtype=params_dtype,
                                                                  quant_config=quant_config,
                                                                  prefix=prefix)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: add self.tp_group, world_size and tp_rank to support moe expert parallel
    '''
    self.tp_group = tp_group
    self.tp_world_size = get_parallel_world_size_with_group(self.tp_group)
    self.tp_rank = get_parallel_rank_with_group(self.tp_group)
    '''
    =================
    End of MLU Hijack
    =================
    '''


def vllm__model_executor__layers__linear__ColumnParallelLinear____init__(
        self,
        input_size: int,
        output_size: int,
        bias: bool = True,
        gather_output: bool = False,
        skip_bias_add: bool = False,
        params_dtype: Optional[torch.dtype] = None,
        quant_config: Optional[QuantizationConfig] = None,
        output_sizes: Optional[List[int]] = None,
        prefix: str = "",
        tp_group: Any = None,
    ):
    super(ColumnParallelLinear, self).__init__(input_size, output_size, skip_bias_add, params_dtype,
                     quant_config, prefix, tp_group)

    self.gather_output = gather_output

    # Divide the weight matrix along the last dimension.
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: modify get_tensor_model_parallel_world_size() to self.tp_world_size
    @brief: move checking output_sizes logic from MergedColumnParallelLinear to here
    '''
    tp_size = self.tp_world_size

    if output_sizes is not None:
        assert all(output_size_var % tp_size == 0 for output_size_var in output_sizes)
    '''
    =================
    End of MLU Hijack
    =================
    '''
    assert self.quant_method is not None
    self.output_size_per_partition = divide(self.output_size, tp_size)
    self.output_partition_sizes = [self.output_size_per_partition]
    # If QKV or MergedColumn, use output size of each partition.
    if hasattr(self, "output_sizes"):
        self.output_partition_sizes = [
            divide(output_size, tp_size)
            for output_size in self.output_sizes
        ]

    if output_sizes is None:
        output_sizes = [output_size]

    self.quant_method.create_weights(
        layer=self,
        input_size_per_partition=self.input_size,
        output_partition_sizes=self.output_partition_sizes,
        input_size=self.input_size,
        output_size=self.output_size,
        params_dtype=self.params_dtype,
        weight_loader=(
            self.weight_loader_v2 if self.quant_method.__class__.__name__
            in WEIGHT_LOADER_V2_SUPPORTED else self.weight_loader))
    if bias:
        self.bias = Parameter(
            torch.empty(self.output_size_per_partition,
                        dtype=params_dtype))
        set_weight_attrs(self.bias, {
            "output_dim": 0,
            "weight_loader": self.weight_loader,
        })
    else:
        self.register_parameter("bias", None)


def vllm__model_executor__layers__linear__ColumnParallelLinear__weight_loader(
        self, param: Parameter, loaded_weight: torch.Tensor):
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: modify get_tensor_model_parallel_rank() to self.tp_rank
    '''
    tp_rank = self.tp_rank
    '''
    =================
    End of MLU Hijack
    =================
    '''
    output_dim = getattr(param, "output_dim", None)

    # Special case for GGUF
    is_gguf_weight = getattr(param, "is_gguf_weight", False)
    is_gguf_weight_type = getattr(param, "is_gguf_weight_type", False)
    if is_gguf_weight_type:
        param.weight_type = loaded_weight.item()

    # Materialize GGUF UninitializedParameter
    if is_gguf_weight and isinstance(param, UninitializedParameter):
        param.materialize(loaded_weight.shape, dtype=loaded_weight.dtype)

    use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False)

    param_data = param.data
    # bitsandbytes loads the weights of the specific portion
    # no need to narrow here
    if output_dim is not None and not use_bitsandbytes_4bit:
        shard_size = param_data.shape[output_dim]
        start_idx = tp_rank * shard_size
        loaded_weight = loaded_weight.narrow(output_dim, start_idx,
                                             shard_size)

    # Special case for loading scales off disk, which often do not
    # have a shape (such as in the case of AutoFP8).
    if len(loaded_weight.shape) == 0:
        loaded_weight = loaded_weight.reshape(1)

    assert param_data.shape == loaded_weight.shape
    param_data.copy_(loaded_weight)


def vllm__model_executor__layers__linear__ColumnParallelLinear__forward(
        self, input_, smooth_quant_scale: Optional[torch.Tensor] = None):
    bias = self.bias if not self.skip_bias_add else None

    # Matrix multiply.
    assert self.quant_method is not None
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: Add input_scale parameter.
    '''
    if smooth_quant_scale is not None:
        output_parallel = self.quant_method.apply(self, input_, bias,
                input_scale=smooth_quant_scale)
    else:
        output_parallel = self.quant_method.apply(self, input_, bias)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    if self.gather_output:
        # All-gather across the partitions.
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: add tp_group param to tensor_model_parallel_all_gather
        '''
        output = tensor_model_parallel_all_gather(output_parallel, self.tp_group)
        '''
        =================
        End of MLU Hijack
        =================
        '''
    else:
        output = output_parallel
    output_bias = self.bias if self.skip_bias_add else None
    return output, output_bias


def vllm__model_executor__layers__linear__ColumnParallelLinear__extra_repr(self) -> str:
    s = f"in_features={self.input_size}"
    s += f", output_features={self.output_size_per_partition}"
    s += f", bias={self.bias is not None}"
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: modify get_tensor_model_parallel_world_size() to self.tp_world_size
    '''
    s += f", tp_size={self.tp_world_size}"
    '''
    =================
    End of MLU Hijack
    =================
    '''
    s += f", gather_output={self.gather_output}"
    return s


def vllm__model_executor__layers__linear__MergedColumnParallelLinear____init__(
        self,
        input_size: int,
        output_sizes: List[int],
        bias: bool = True,
        gather_output: bool = False,
        skip_bias_add: bool = False,
        params_dtype: Optional[torch.dtype] = None,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        tp_group: Any = None,
    ):
    self.output_sizes = output_sizes
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: move checking output_sizes logic from MergedColumnParallelLinear to ColumnParallelLinear.__init__
    '''
    # tp_size = get_tensor_model_parallel_world_size()
    # assert all(output_size % tp_size == 0 for output_size in output_sizes)
    '''
    =================
    End of MLU Hijack
    =================
    '''
    super(MergedColumnParallelLinear, self).__init__(input_size=input_size,
                                                     output_size=sum(output_sizes),
                                                     bias=bias,
                                                     gather_output=gather_output,
                                                     skip_bias_add=skip_bias_add,
                                                     params_dtype=params_dtype,
                                                     quant_config=quant_config,
                                                     output_sizes=self.output_sizes,
                                                     prefix=prefix,
                                                     tp_group=tp_group)


def vllm__model_executor__layers__linear__MergedColumnParallelLinear__weight_loader(self,
                      param: Parameter,
                      loaded_weight: torch.Tensor,
                      loaded_shard_id: Optional[int] = None):
    # Special case for GGUF
    # initialize GGUF param after we know the quantize type
    is_gguf_weight = getattr(param, "is_gguf_weight", False)
    is_gguf_weight_type = getattr(param, "is_gguf_weight_type", False)
    if is_gguf_weight_type:
        param.data[loaded_shard_id].copy_(loaded_weight)
        param.shard_weight_type[loaded_shard_id] = loaded_weight.item()
        return

    if is_gguf_weight:
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: modify get_tensor_model_parallel_rank() to self.tp_rank
        @brief: modify get_tensor_model_parallel_world_size() to self.tp_world_size
        '''
        tp_rank = self.tp_rank
        tp_size = self.tp_world_size
        '''
        =================
        End of MLU Hijack
        =================
        '''
        output_dim = getattr(param, "output_dim", None)
        shard_size = loaded_weight.size(output_dim) // tp_size
        start_idx = tp_rank * shard_size

        loaded_weight = loaded_weight.narrow(output_dim, start_idx,
                                             shard_size)

        param.shard_id.append(loaded_shard_id)
        param.shard_id_map[loaded_shard_id] = len(param.data_container)
        param.data_container.append(loaded_weight)
        if len(param.data_container) == 2:
            self.qweight = param.materialize_nested()
        return

    param_data = param.data
    output_dim = getattr(param, "output_dim", None)
    # Special case for AQLM codebooks.
    is_metadata = getattr(param, "is_metadata", False)
    # Special case for per-tensor scale to load scalar into fused array.
    needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False)

    if loaded_shard_id is None:
        # Loaded weight is already fused on disk (qkv/mlp).
        if output_dim is None:
            if needs_scalar_to_array:
                param_data, loaded_weight = adjust_scalar_to_fused_array(
                    param_data, loaded_weight, 0)

            assert param_data.shape == loaded_weight.shape
            param_data.copy_(loaded_weight)
            return
        current_shard_offset = 0
        shard_offsets: List[Tuple[int, int, int]] = []
        for i, output_size in enumerate(self.output_sizes):
            shard_offsets.append((i, current_shard_offset, output_size))
            current_shard_offset += output_size
        packed_dim = getattr(param, "packed_dim", None)
        for shard_id, shard_offset, shard_size in shard_offsets:
            # Special case for Quantization.
            # If quantized, we need to adjust the offset and size to account
            # for the packing.
            if packed_dim == output_dim:
                shard_size = shard_size // param.pack_factor
                shard_offset = shard_offset // param.pack_factor
                # Special case for Marlin.
                shard_size, shard_offset = adjust_marlin_shard(
                    param, shard_size, shard_offset)

            loaded_weight_shard = loaded_weight.narrow(
                output_dim, shard_offset, shard_size)
            self.weight_loader(param, loaded_weight_shard, shard_id)
        return

    assert loaded_shard_id < len(self.output_sizes)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: modify get_tensor_model_parallel_rank() to self.tp_rank
    @brief: modify get_tensor_model_parallel_world_size() to self.tp_world_size
    '''
    tp_rank = self.tp_rank
    tp_size = self.tp_world_size
    '''
    =================
    End of MLU Hijack
    =================
    '''
    if output_dim is not None:
        shard_offset = sum(self.output_sizes[:loaded_shard_id]) // tp_size
        shard_size = self.output_sizes[loaded_shard_id] // tp_size
        # Special case for quantization.
        # If quantized, we need to adjust the offset and size to account
        # for the packing.
        packed_dim = getattr(param, "packed_dim", None)
        if packed_dim == output_dim:
            shard_size = shard_size // param.pack_factor
            shard_offset = shard_offset // param.pack_factor
            # Special case for Marlin.
            shard_size, shard_offset = adjust_marlin_shard(
                param, shard_size, shard_offset)

        use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit",
                                        False)
        if use_bitsandbytes_4bit:
            shard_size = loaded_weight.shape[output_dim]
            shard_offset = loaded_weight.shape[output_dim] * \
                loaded_shard_id

        param_data = param_data.narrow(output_dim, shard_offset,
                                       shard_size)
        start_idx = tp_rank * shard_size
        # bitsandbytes loads the weights of the specific portion
        # no need to narrow here
        if not use_bitsandbytes_4bit:
            loaded_weight = loaded_weight.narrow(output_dim, start_idx,
                                                 shard_size)
    # Special case for AQLM codebooks.
    elif is_metadata:
        # metadata indicates fixed size concatenated along dim 0
        shard_size = loaded_weight.shape[0]
        shard_offset = loaded_shard_id * shard_size
        param_data = param_data.narrow(0, shard_offset, shard_size)

    # Special case for per-tensor scales in fused case.
    elif needs_scalar_to_array:
        param_data, loaded_weight = adjust_scalar_to_fused_array(
            param_data, loaded_weight, loaded_shard_id)

    else:
        ignore_warning = getattr(param, "ignore_warning", False)
        if not ignore_warning:
            logger.warning(
                "Loading a weight without `output_dim` attribute in "
                "MergedColumnParallelLinear, assume the weight is "
                "the same for all partitions.")

    assert param_data.shape == loaded_weight.shape
    param_data.copy_(loaded_weight)


def vllm__model_executor__layers__linear__MergedColumnParallelLinear__weight_loader_v2(self,
                         param: BasevLLMParameter,
                         loaded_weight: torch.Tensor,
                         loaded_shard_id: Optional[int] = None):
    if loaded_shard_id is None:
        if isinstance(param, PerTensorScaleParameter):
            param.load_merged_column_weight(loaded_weight=loaded_weight,
                                            shard_id=0)
            return
        elif type(param) in (RowvLLMParameter, BasevLLMParameter):
            param.load_merged_column_weight(loaded_weight=loaded_weight)
            return
        # TODO: @dsikka - move to parameter.py
        self._load_fused_module_from_checkpoint(param, loaded_weight)
        return

    assert loaded_shard_id < len(self.output_sizes)

    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: modify get_tensor_model_parallel_world_size() to self.tp_world_size
    '''
    tp_size = self.tp_world_size
    '''
    =================
    End of MLU Hijack
    =================
    '''
    shard_offset = sum(self.output_sizes[:loaded_shard_id]) // tp_size
    shard_size = self.output_sizes[loaded_shard_id] // tp_size

    param.load_merged_column_weight(loaded_weight=loaded_weight,
                                    shard_id=loaded_shard_id,
                                    shard_offset=shard_offset,
                                    shard_size=shard_size)

def vllm__model_executor__layers__linear__RowParallelLinear____init__(
        self,
        input_size: int,
        output_size: int,
        bias: bool = True,
        input_is_parallel: bool = True,
        skip_bias_add: bool = False,
        params_dtype: Optional[torch.dtype] = None,
        reduce_results: bool = True,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
        tp_group: Any = None,
    ):
    super(RowParallelLinear, self).__init__(input_size, output_size, skip_bias_add, params_dtype,
                     quant_config, prefix, tp_group)

    self.input_is_parallel = input_is_parallel
    self.reduce_results = reduce_results

    # Divide the weight matrix along the last dimension.
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: modify get_tensor_model_parallel_rank() to self.tp_rank
    @brief: modify get_tensor_model_parallel_world_size() to self.tp_world_size
    '''
    self.tp_size = self.tp_world_size
    '''
    =================
    End of MLU Hijack
    =================
    '''
    self.input_size_per_partition = divide(input_size, self.tp_size)
    assert self.quant_method is not None

    self.quant_method.create_weights(
        layer=self,
        input_size_per_partition=self.input_size_per_partition,
        output_partition_sizes=[self.output_size],
        input_size=self.input_size,
        output_size=self.output_size,
        params_dtype=self.params_dtype,
        weight_loader=(
            self.weight_loader_v2 if self.quant_method.__class__.__name__
            in WEIGHT_LOADER_V2_SUPPORTED else self.weight_loader))
    if not reduce_results and (bias and not skip_bias_add):
        raise ValueError("When not reduce the results, adding bias to the "
                         "results can lead to incorrect results")

    if bias:
        self.bias = Parameter(
            torch.empty(self.output_size, dtype=params_dtype))
        set_weight_attrs(self.bias, {
            "output_dim": 0,
            "weight_loader": self.weight_loader,
        })
    else:
        self.register_parameter("bias", None)


def vllm__model_executor__layers__linear__RowParallelLinear__weight_loader(
        self, param: Parameter, loaded_weight: torch.Tensor):
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: modify get_tensor_model_parallel_rank() to self.tp_rank
    @brief: modify get_tensor_model_parallel_world_size() to self.tp_world_size
    '''
    tp_rank = self.tp_rank
    tp_size = self.tp_world_size
    '''
    =================
    End of MLU Hijack
    =================
    '''
    input_dim = getattr(param, "input_dim", None)
    use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False)

    # Special case for GGUF
    is_gguf_weight = getattr(param, "is_gguf_weight", False)
    is_gguf_weight_type = getattr(param, "is_gguf_weight_type", False)
    if is_gguf_weight_type:
        param.weight_type = loaded_weight.item()

    # Materialize GGUF UninitializedParameter
    if is_gguf_weight and isinstance(param, UninitializedParameter):
        weight_shape = list(loaded_weight.shape)
        if input_dim:
            weight_shape[input_dim] = weight_shape[input_dim] // tp_size
        param.materialize(tuple(weight_shape), dtype=loaded_weight.dtype)

    param_data = param.data
    # bitsandbytes loads the weights of the specific portion
    # no need to narrow here
    if input_dim is not None and not use_bitsandbytes_4bit:
        shard_size = param_data.shape[input_dim]
        start_idx = tp_rank * shard_size
        loaded_weight = loaded_weight.narrow(input_dim, start_idx,
                                             shard_size)

    # Special case for loading scales off disk, which often do not
    # have a shape (such as in the case of AutoFP8).
    if len(loaded_weight.shape) == 0:
        loaded_weight = loaded_weight.reshape(1)

    assert param_data.shape == loaded_weight.shape
    param_data.copy_(loaded_weight)


def vllm__model_executor__layers__linear__RowParallelLinear__forward(
    self,
    input_,
    residual: Optional[torch.Tensor] = None
):
    if self.input_is_parallel:
        input_parallel = input_
    else:
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: modify get_tensor_model_parallel_rank() to self.tp_rank
        '''
        tp_rank = self.tp_rank
        '''
        =================
        End of MLU Hijack
        =================
        '''
        splitted_input = split_tensor_along_last_dim(
            input_, num_partitions=self.tp_size)
        input_parallel = splitted_input[tp_rank].contiguous()

    # Matrix multiply.
    assert self.quant_method is not None
    # Only fuse bias add into GEMM for rank 0 (this ensures that
    # bias will not get added more than once in TP>1 case)
    bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
    residual_ = None if self.tp_rank > 0 else residual
    '''
    =====================================================
    Modify by custom vllm_mlu
    =====================================================
    @brief: abandon original reduce if parallel_num is set
    '''
    is_parallel_enable = hasattr(self.quant_method, 'parallel_num') and get_is_prompt()
    '''
    =====================================================
    End of custom MLU Hijack
    =====================================================
    '''
    output_parallel = self.quant_method.apply(self,
                                              input_parallel,
                                              bias=bias_,
                                              residual=residual_)
    '''
    =============================
    Modify by custom vllm_mlu
    =============================
    @brief: when preload_size is set, call GroupCoordinator.all_reduce() directly and
    use async_op to set all_reduce paralleled with preload
    '''
    if self.reduce_results and self.tp_size > 1 and not is_parallel_enable:
        if hasattr(self, 'preload_size') and self.preload_size > 0 and not self.is_prompt:
            handle = get_tp_group(self.tp_group).all_reduce(output_parallel, async_op=True)
            _MB = 1 << 20
            mlu_ops.preload(self.preloaded_weights[0].data, self.preload_size * _MB)
            preloaded_weights_size = self.preloaded_weights[0].numel() * self.preloaded_weights[0].element_size()
            if preloaded_weights_size < (self.preload_size * _MB) and len(self.preloaded_weights) > 1:
                mlu_ops.preload(self.preloaded_weights[1].data, (self.preload_size * _MB) - preloaded_weights_size)
            handle.wait()
            output = output_parallel
        else:
            '''
             =============================
             Modify by vllm_mlu
             =============================
             @brief: add tensor_model_parallel_all_reduce() with self.tp_group
             '''
            output = tensor_model_parallel_all_reduce(output_parallel, tp_group=self.tp_group)
            '''
             =================
             End of MLU Hijack
             =================
             '''
    else:
        output = output_parallel
    '''
    =========================
    End of custom MLU Hijack
    =========================
    '''
    output_bias = self.bias if self.skip_bias_add else None

    return output, output_bias


MluHijackObject.apply_hijack(LinearBase,
                             LinearBase.__init__,
                             vllm__model_executor__layers__linear__LinearBase____init__)
MluHijackObject.apply_hijack(ColumnParallelLinear,
                             ColumnParallelLinear.__init__,
                             vllm__model_executor__layers__linear__ColumnParallelLinear____init__)
MluHijackObject.apply_hijack(ColumnParallelLinear,
                             ColumnParallelLinear.weight_loader,
                             vllm__model_executor__layers__linear__ColumnParallelLinear__weight_loader)
MluHijackObject.apply_hijack(ColumnParallelLinear,
                             ColumnParallelLinear.forward,
                             vllm__model_executor__layers__linear__ColumnParallelLinear__forward)
MluHijackObject.apply_hijack(ColumnParallelLinear,
                             ColumnParallelLinear.extra_repr,
                             vllm__model_executor__layers__linear__ColumnParallelLinear__extra_repr)
MluHijackObject.apply_hijack(MergedColumnParallelLinear,
                             MergedColumnParallelLinear.__init__,
                             vllm__model_executor__layers__linear__MergedColumnParallelLinear____init__)
MluHijackObject.apply_hijack(MergedColumnParallelLinear,
                             MergedColumnParallelLinear.weight_loader,
                             vllm__model_executor__layers__linear__MergedColumnParallelLinear__weight_loader)
MluHijackObject.apply_hijack(MergedColumnParallelLinear,
                             MergedColumnParallelLinear.weight_loader_v2,
                             vllm__model_executor__layers__linear__MergedColumnParallelLinear__weight_loader_v2)
MluHijackObject.apply_hijack(RowParallelLinear,
                             RowParallelLinear.__init__,
                             vllm__model_executor__layers__linear__RowParallelLinear____init__)
MluHijackObject.apply_hijack(RowParallelLinear,
                             RowParallelLinear.weight_loader,
                             vllm__model_executor__layers__linear__RowParallelLinear__weight_loader)
MluHijackObject.apply_hijack(RowParallelLinear,
                             RowParallelLinear.forward,
                             vllm__model_executor__layers__linear__RowParallelLinear__forward)
