from typing import Optional

import torch

from vllm.distributed import tensor_model_parallel_all_reduce
from vllm.lora.fully_sharded_layers import RowParallelLinearWithShardedLoRA
from vllm_mlu.mlu_hijack_utils import MluHijackObject


def vllm__lora__fully_sharded_layers__RowParallelLinearWithShardedLoRA__apply(
    self,
    x: torch.Tensor,
    bias: Optional[torch.Tensor],
    residual: Optional[torch.Tensor]
) -> torch.Tensor:
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: add residual and bias in matmul
    '''
    output = self.base_layer.quant_method.apply(
        self.base_layer, x, bias, residual)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    x = x.view(-1, x.shape[-1])
    output, out_orig_shape = output.view(-1,
                                         output.shape[-1]), output.shape
    buffer = torch.zeros(
        (x.shape[0], self.lora_a_stacked.shape[2]),
        dtype=torch.float32,
        device=x.device,
    )

    self.punica_wrapper.add_shrink(buffer, x, self.lora_a_stacked, 1.0)
    buffer = tensor_model_parallel_all_reduce(buffer)

    # following S-LoRA, allows the fusing of all_gather and all_reduce
    # by adding the column partitioned lora output to a slice of output
    # tensor, which is a partial sum due to row parallel. All that
    # remains is a standard all_reduce. User should be aware though that
    # the output is not the same as a normal row_parallel, it should be
    # reduced before being used
    shard_size = self.lora_b_stacked.shape[2]
    start_idx = self.tp_rank * shard_size

    if self.bias_stacked is not None:
        bias = self.bias_stacked.view(-1, self.bias_stacked.shape[-1])
        bias = bias[self.punica_wrapper.token_lora_indices]
        bias[self.punica_wrapper.token_lora_indices == -1] = 0
        output += bias

    self.punica_wrapper.add_expand_slice(output, buffer,
                                         self.lora_b_stacked, start_idx,
                                         shard_size)
    output = output.view(*out_orig_shape)
    return output


MluHijackObject.apply_hijack(RowParallelLinearWithShardedLoRA,
                             RowParallelLinearWithShardedLoRA.apply,
                             vllm__lora__fully_sharded_layers__RowParallelLinearWithShardedLoRA__apply)
