
import torch
from torch import nn
from typing import Any, Dict, Iterable, List, Optional, Tuple
from transformers import PretrainedConfig

from vllm.attention import Attention, AttentionMetadata
from vllm.config import CacheConfig
from vllm.distributed import (get_tensor_model_parallel_world_size,
                              tensor_model_parallel_all_reduce)
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.fused_moe import FusedMoE
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
                                               ReplicatedLinear,
                                               RowParallelLinear)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig)
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.layers.sampler import Sampler
from vllm.model_executor.layers.vocab_parallel_embedding import (
    ParallelLMHead, VocabParallelEmbedding)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.sampling_metadata import SamplingMetadata
from vllm.sequence import IntermediateTensors
from vllm.model_executor.layers.sampler import SamplerOutput
from vllm_mlu.model_executor.layers.feed_forward import FeedForward
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp
from vllm.utils import print_warning_once
from vllm.model_executor.models.utils import is_pp_missing_parameter
from vllm_mlu.model_executor.models.layer_utils import quant_fusion_with_rmsnorm

from vllm.model_executor.models.deepseek_v2 import DeepseekV2Attention, DeepseekV2DecoderLayer, DeepseekV2ForCausalLM, yarn_get_mscale

class DeepseekV2MoE(SparseMoeMlp):

    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
        prefix: str = "",
    ):
        super().__init__(num_experts=config.n_routed_experts,
                         top_k=config.num_experts_per_tok,
                         hidden_size=config.hidden_size,
                         intermediate_size=config.moe_intermediate_size,
                         up_proj_name="gate_up_proj",
                         is_gated=True,
                         down_proj_name="down_proj",
                         has_bias=False,
                         skip_bias_add=False,
                         renormalize=config.norm_topk_prob,
                         hidden_act=config.hidden_act,
                         params_dtype=None,
                         quant_config=quant_config,
                         is_use_fused_moe=True,
                         expert_group=config.n_group,
                         topk_group=config.topk_group)
        self.config = config
        self.tp_size = get_tensor_model_parallel_world_size()
        self.routed_scaling_factor = config.routed_scaling_factor
        self.n_shared_experts = config.n_shared_experts
        self.routed_scaling_factor = config.routed_scaling_factor
        if self.tp_size > config.n_routed_experts:
            raise ValueError(
                f"Tensor parallel size {self.tp_size} is greater than "
                f"the number of experts {config.n_routed_experts}.")

        if config.hidden_act != "silu":
            raise ValueError(f"Unsupported activation: {config.hidden_act}. "
                             "Only silu is supported for now.")

        self.gate = ReplicatedLinear(config.hidden_size,
                                     config.n_routed_experts,
                                     bias=False,
                                     quant_config=None,
                                     prefix=f"{prefix}.gate")
        if config.n_shared_experts is not None:
            intermediate_size = (config.moe_intermediate_size *
                                 config.n_shared_experts)
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: replace MLP with FeedForward.
            '''
            self.shared_experts = FeedForward(hidden_size=config.hidden_size,
                                             intermediate_size=intermediate_size,
                                             hidden_act=config.hidden_act,
                                             up_proj_name='gate_up_proj',
                                             is_gated=True,
                                             down_proj_name='down_proj',
                                             bias=False,
                                             quant_config=quant_config,
                                             reduce_results=False)
            '''
            ==================
            End of MLU Hijack
            ==================
            '''

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        num_tokens, hidden_dim = hidden_states.shape
        hidden_states = hidden_states.view(-1, hidden_dim)
        if self.n_shared_experts is not None:
            shared_output = self.shared_experts(hidden_states)
        # router_logits: (num_tokens, n_experts)
        router_logits, _ = self.gate(hidden_states)
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: replace experts() with forward_experts, which defined by SparseMoeMlp.
        '''
        final_hidden_states = self.forward_experts(
            hidden_states, router_logits) * self.routed_scaling_factor
        '''
        ==================
        End of MLU Hijack
        ==================
        '''
        if shared_output is not None:
            final_hidden_states = final_hidden_states + shared_output
        if self.tp_size > 1:
            final_hidden_states = tensor_model_parallel_all_reduce(
                final_hidden_states)

        return final_hidden_states.view(num_tokens, hidden_dim)

def forward_prefill(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
    kv_cache: torch.Tensor,
    attn_metadata: AttentionMetadata,
) -> torch.Tensor:
    if self.q_lora_rank is not None:
        q = self.q_a_proj(hidden_states)[0]
        q_scale = None
        if hasattr(self.q_b_proj.quant_method, "quant_config"):
            self.q_b_proj.quant_method.skip_quant_input = True
            quant_scale = self.q_b_proj.smooth
            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
                        self.q_a_layernorm, quant_scale, dynamic_quant=True)
            q, q_scale = self.quant_fusion_attn_layernorm(q)
            q = self.q_b_proj(q, q_scale)[0].view(-1, self.num_local_heads, self.qk_head_dim)
        else:
            q = self.q_a_layernorm(q)
            q = self.q_b_proj(q)[0].view(-1, self.num_local_heads,
                                         self.qk_head_dim)
    else:
        q = self.q_proj(hidden_states)[0].view(-1, self.num_local_heads,
                                               self.qk_head_dim)
    q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim],
                           dim=-1)
    latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
    kv_a, _ = latent_cache.split(
        [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
    latent_cache = latent_cache.unsqueeze(1)
    kv_a = self.kv_a_layernorm(kv_a.contiguous())
    kv = self.kv_b_proj(kv_a)[0]
    kv = kv.view(-1, self.num_local_heads,
                 self.qk_nope_head_dim + self.v_head_dim)
    k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
    k_pe = latent_cache[:, :, self.kv_lora_rank:]
    q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: MLA save cache before flashattn
    '''
    from vllm import _mlu_ops as mlu_ops
    if len(kv_cache) != 0  and kv_cache[0].numel() > 0:
        key_cache = kv_cache[0][0]
        key_value = torch.concat((kv_a.unsqueeze(1), k_pe), dim=-1)
        updated_slot_mapping = attn_metadata.slot_mapping
        if self.attn.kv_cache_dtype == 'int8':
            key_cache_scale = kv_cache[1][0]
            mlu_ops.quant_to_paged_cache(key_value,
                                         None,
                                         key_cache,
                                         None,
                                         key_cache_scale,
                                         None,
                                         updated_slot_mapping.flatten())
        else:
            mlu_ops.reshape_paged_cache(key_value,
                                        None,
                                        key_cache,
                                        None,
                                        updated_slot_mapping.flatten())
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    q[..., self.qk_nope_head_dim:] = q_pe
    k = torch.empty_like(q)
    k[..., :self.qk_nope_head_dim] = k_nope
    k[..., self.qk_nope_head_dim:] = k_pe
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: mlu attention not pad but qk_head_dim 192 v_head_dim 128.
    '''
    q = q.reshape(-1, self.num_local_heads * self.qk_head_dim)
    k = k.reshape(-1, self.num_local_heads * self.qk_head_dim)
    v = v.contiguous().reshape(-1, self.num_local_heads * self.v_head_dim)

    attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
    attn_output = attn_output.reshape(-1, self.num_local_heads * self.v_head_dim)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    output, _ = self.o_proj(attn_output)
    return output

def forward_decoder(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
    kv_cache: torch.Tensor,
    attn_metadata: AttentionMetadata,
) -> torch.Tensor:
    q_len = hidden_states.shape[0]
    q_input = hidden_states.new_empty(
        q_len, self.num_local_heads, self.kv_lora_rank + self.qk_rope_head_dim
    )
    if self.q_lora_rank is not None:
        q = self.q_a_proj(hidden_states)[0]
        q_scale = None
        if hasattr(self.q_b_proj.quant_method, "quant_config"):
            self.q_b_proj.quant_method.skip_quant_input = True
            quant_scale = self.q_b_proj.smooth
            self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
                        self.q_a_layernorm, quant_scale, dynamic_quant=True)
            q, q_scale = self.quant_fusion_attn_layernorm(q)
            q = self.q_b_proj(q, q_scale)[0].view(-1, self.num_local_heads, self.qk_head_dim)
        else:
            q = self.q_a_layernorm(q)
            q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
    else:
        q = self.q_proj(hidden_states)[0].view(
            -1, self.num_local_heads, self.qk_head_dim
        )
    q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

    from vllm import _mlu_ops as mlu_ops
    q_nope_out = torch.bmm(q_nope.transpose(0, 1), self.w_kc)
    q_input[..., : self.kv_lora_rank] = q_nope_out.transpose(0, 1)

    latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
    v_input = latent_cache[..., : self.kv_lora_rank]
    v_input = self.kv_a_layernorm(v_input.contiguous()).unsqueeze(1)
    k_input = latent_cache.unsqueeze(1)
    k_input[..., : self.kv_lora_rank] = v_input
    k_pe = k_input[..., self.kv_lora_rank :]

    q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
    q_input[..., self.kv_lora_rank :] = q_pe
    k_input[..., self.kv_lora_rank :] = k_pe
    v_input = torch.nn.functional.pad(v_input, [0, self.qk_rope_head_dim, 0, 0, 0, 0],
                                value=0).view(-1, self.kv_lora_rank + self.qk_rope_head_dim)
    q_input = q_input.reshape(q_input.shape[0], -1)
    k_input = k_input.reshape(k_input.shape[0], -1)
    v_input = v_input.reshape(v_input.shape[0], -1)
    attn_output = self.attn_decoder(q_input, k_input, v_input, kv_cache, attn_metadata)
    attn_output = attn_output.reshape(-1, self.num_local_heads,
                                      self.kv_lora_rank + self.qk_rope_head_dim)
    attn_output = attn_output[:, :, :self.kv_lora_rank]
    attn_bmm_output = torch.bmm(attn_output.transpose(0, 1), self.w_vc.transpose(1, 2).contiguous())
    attn_output = attn_bmm_output.transpose(0, 1).flatten(1, 2)
    output, _ = self.o_proj(attn_output)
    return output

def vllm__module_executor__models__deepseek_v2__DeepseekV2Attention__init(
    self,
    config: PretrainedConfig,
    hidden_size: int,
    num_heads: int,
    qk_nope_head_dim: int,
    qk_rope_head_dim: int,
    v_head_dim: int,
    q_lora_rank: int,
    kv_lora_rank: int,
    rope_theta: float = 10000,
    rope_scaling: Optional[Dict[str, Any]] = None,
    max_position_embeddings: int = 8192,
    cache_config: Optional[CacheConfig] = None,
    quant_config: Optional[QuantizationConfig] = None,
    prefix: str = "",
) -> None:
    super(DeepseekV2Attention, self).__init__()
    self.hidden_size = hidden_size
    self.qk_nope_head_dim = qk_nope_head_dim
    self.qk_rope_head_dim = qk_rope_head_dim
    self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
    self.v_head_dim = v_head_dim
    self.q_lora_rank = q_lora_rank
    self.kv_lora_rank = kv_lora_rank
    self.num_heads = num_heads
    tp_size = get_tensor_model_parallel_world_size()
    assert num_heads % tp_size == 0
    self.num_local_heads = num_heads // tp_size
    self.scaling = self.qk_head_dim**-0.5
    self.rope_theta = rope_theta
    self.max_position_embeddings = max_position_embeddings

    if self.q_lora_rank is not None:
        # only RowParallelLinear/ColumnParallelLinear will be quantize
        self.q_a_proj = ReplicatedLinear(self.hidden_size,
                                         self.q_lora_rank,
                                         bias=False,
                                         quant_config=None,
                                         prefix=f"{prefix}.q_a_proj")
        self.q_a_layernorm = RMSNorm(self.q_lora_rank,
                                     eps=config.rms_norm_eps)
        self.q_b_proj = ColumnParallelLinear(q_lora_rank,
                                             self.num_heads *
                                             self.qk_head_dim,
                                             bias=False,
                                             quant_config=quant_config,
                                             prefix=f"{prefix}.q_b_proj")
    else:
        self.q_proj = ColumnParallelLinear(self.hidden_size,
                                           self.num_heads *
                                           self.qk_head_dim,
                                           bias=False,
                                           quant_config=quant_config,
                                           prefix=f"{prefix}.q_proj")

    self.kv_a_proj_with_mqa = ReplicatedLinear(
        self.hidden_size,
        self.kv_lora_rank + self.qk_rope_head_dim,
        bias=False,
        quant_config=None,
        prefix=f"{prefix}.kv_a_proj_with_mqa")
    self.kv_a_layernorm = RMSNorm(self.kv_lora_rank,
                                  eps=config.rms_norm_eps)
    self.kv_b_proj = ColumnParallelLinear(
        self.kv_lora_rank,
        self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
        bias=False,
        quant_config=None,
        prefix=f"{prefix}.kv_b_proj")
    kv_b_proj_weight = self.kv_b_proj.weight
    w_kc, w_vc = kv_b_proj_weight.unflatten(
        0, (-1, self.qk_nope_head_dim + self.v_head_dim)
        ).split([self.qk_nope_head_dim, self.v_head_dim], dim=1)
    self.w_kc = w_kc
    self.w_vc = w_vc
    # O projection.
    self.o_proj = RowParallelLinear(self.num_heads * self.v_head_dim,
                                    self.hidden_size,
                                    bias=False,
                                    quant_config=quant_config,
                                    prefix=f"{prefix}.o_proj")
    rope_scaling['rope_type'] = 'deepseek_yarn'
    self.rotary_emb = get_rope(qk_rope_head_dim,
                               rotary_dim=qk_rope_head_dim,
                               max_position=max_position_embeddings,
                               base=rope_theta,
                               rope_scaling=rope_scaling,
                               is_neox_style=False)

    if rope_scaling:
        mscale_all_dim = rope_scaling.get("mscale_all_dim", False)
        scaling_factor = rope_scaling["factor"]
        mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
        self.scaling = self.scaling * mscale * mscale

    # self.attn = Attention(self.num_heads,
    #                       self.qk_head_dim,
    #                       self.scaling,
    #                       num_kv_heads=self.num_heads)

    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: mlu attention support head_size 192.
    '''
    self.attn = Attention(self.num_local_heads,
                          self.qk_head_dim,
                          self.scaling,
                          num_kv_heads=self.num_local_heads,
                          cache_config=cache_config,
                          quant_config=quant_config,
                          use_mla=True)
    self.attn_decoder = Attention(self.num_local_heads,
                          self.kv_lora_rank + self.qk_rope_head_dim,
                          self.scaling,
                          num_kv_heads=1,
                          cache_config=cache_config,
                          quant_config=quant_config,
                          use_mla=True)
    import types
    self.forward_prefill = types.MethodType(forward_prefill, self)
    self.forward_decoder = types.MethodType(forward_decoder, self)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''

def vllm__module_executor__models__deepseek_v2__DeepseekV2Attention__forward(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
    kv_cache: torch.Tensor,
    attn_metadata: AttentionMetadata,
) -> torch.Tensor:
    # Use normal computation for prefill and use weight absorption for extend/decode
    if attn_metadata.prefill_metadata:
        return self.forward_prefill(positions, hidden_states, kv_cache,
                                    attn_metadata)
    else:
        return self.forward_decoder(positions, hidden_states, kv_cache,
                                    attn_metadata)


def vllm__module_executor__models__deepseek_v2__DeepseekV2DecoderLayer__init__(
    self,
    config: PretrainedConfig,
    prefix: str,
    cache_config: Optional[CacheConfig] = None,
    quant_config: Optional[QuantizationConfig] = None,
) -> None:
    super(DeepseekV2DecoderLayer, self).__init__()
    self.hidden_size = config.hidden_size
    rope_theta = getattr(config, "rope_theta", 10000)
    rope_scaling = getattr(config, "rope_scaling", None)
    max_position_embeddings = getattr(config, "max_position_embeddings",
                                      8192)
    # DecoderLayers are created with `make_layers` which passes the prefix
    # with the layer's index.
    layer_idx = int(prefix.split(sep='.')[-1])
    self.self_attn = DeepseekV2Attention(
        config=config,
        hidden_size=self.hidden_size,
        num_heads=config.num_attention_heads,
        qk_nope_head_dim=config.qk_nope_head_dim,
        qk_rope_head_dim=config.qk_rope_head_dim,
        v_head_dim=config.v_head_dim,
        q_lora_rank=config.q_lora_rank
        if hasattr(config, "q_lora_rank") else None,
        kv_lora_rank=config.kv_lora_rank,
        rope_theta=rope_theta,
        rope_scaling=rope_scaling,
        max_position_embeddings=max_position_embeddings,
        cache_config=cache_config,
        quant_config=quant_config,
        prefix=f"{prefix}.self_attn",
    )
    if (config.n_routed_experts is not None
            and layer_idx >= config.first_k_dense_replace
            and layer_idx % config.moe_layer_freq == 0):
        self.mlp = DeepseekV2MoE(
            config=config,
            quant_config=quant_config,
            prefix=f"{prefix}.mlp",
        )
    else:
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: replace MLP with FeedForward.
        '''
        self.mlp = FeedForward(hidden_size=config.hidden_size,
                               intermediate_size=config.intermediate_size,
                               hidden_act=config.hidden_act,
                               up_proj_name='gate_up_proj',
                               is_gated=True,
                               down_proj_name='down_proj',
                               bias=False,
                               quant_config=quant_config)
        '''
        ==================
        End of MLU Hijack
        ==================
        '''
    self.input_layernorm = RMSNorm(config.hidden_size,
                                   eps=config.rms_norm_eps)
    self.post_attention_layernorm = RMSNorm(config.hidden_size,
                                            eps=config.rms_norm_eps)

def vllm__module_executor__models__deepseek_v2__DeepseekV2ForCausalLM__load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: pack params and cal start expert id
    '''
    for name, m in self.model.named_modules():
        if isinstance(m, SparseMoeMlp):
            m.pack_params()

    start_expert_id = 0
    '''
    ==================
    End of MLU Hijack
    ==================
    '''

    stacked_params_mapping = [
        # (param_name, shard_name, shard_id)
        ("gate_up_proj", "gate_proj", 0),
        ("gate_up_proj", "up_proj", 1),
    ]

    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: delete expert_params_mapping for no useless
    '''
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    params_dict = dict(self.named_parameters())
    for name, loaded_weight in weights:
        if "rotary_emb.inv_freq" in name:
            continue
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: replace expert_id in weight to named_expert_id in params_dict
        '''
        if start_expert_id > 0 and "mlp.experts." in name:
            expert_str = re.search(r'experts\.\d+', name).group(0)
            expert_id=int(expert_str.split(".")[1])
            named_expert_id = expert_id - start_expert_id
            old_expert_name = f"experts.{expert_id}"
            new_expert_name = f"experts.{named_expert_id}"
            name = name.replace(old_expert_name, new_expert_name)
        '''
        ==================
        End of MLU Hijack
        ==================
        '''
        for (param_name, weight_name, shard_id) in stacked_params_mapping:
            # Skip non-stacked layers and experts (experts handled below).
            if weight_name not in name:
                continue
            # We have mlp.experts[0].gate_proj in the checkpoint.
            # Since we handle the experts below in expert_params_mapping,
            # we need to skip here BEFORE we update the name, otherwise
            # name will be updated to mlp.experts[0].gate_up_proj, which
            # will then be updated below in expert_params_mapping
            # for mlp.experts[0].gate_gate_up_proj, which breaks load.
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: add expert skiped condition and delete useless if name not in params_dict: continue condition
            '''
            name = name.replace(weight_name, param_name)
            if (("mlp.experts." in name or "mlp.shared_experts." in name or "mlp.shared_expert_gate." in name)
                    and name not in params_dict):
                continue
            '''
            ==================
            End of MLU Hijack
            ==================
            '''
            # Skip loading extra bias for GPTQ models.
            if name.endswith(".bias") and name not in params_dict:
                continue

            if is_pp_missing_parameter(name, self):
                continue

            param = params_dict[name]
            weight_loader = param.weight_loader
            weight_loader(param, loaded_weight, shard_id)
            break
        else:
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: add expert skiped condition
            '''
            # Skip loading extra bias for GPTQ models.
            if name.endswith(".bias") and name not in params_dict:
                continue

            if (("mlp.experts." in name or "mlp.shared_experts." in name or "mlp.shared_expert_gate." in name)
                    and name not in params_dict):
                continue
            if is_pp_missing_parameter(name, self):
                continue

            param = params_dict[name]
            weight_loader = getattr(param, "weight_loader",
                                    default_weight_loader)
            weight_loader(param, loaded_weight)
            '''
            ==================
            End of MLU Hijack
            ==================
            '''

MluHijackObject.apply_hijack(DeepseekV2Attention,
                             DeepseekV2Attention.forward,
                             vllm__module_executor__models__deepseek_v2__DeepseekV2Attention__forward)
MluHijackObject.apply_hijack(DeepseekV2Attention,
                             DeepseekV2Attention.__init__,
                             vllm__module_executor__models__deepseek_v2__DeepseekV2Attention__init)
MluHijackObject.apply_hijack(DeepseekV2DecoderLayer,
                             DeepseekV2DecoderLayer.__init__,
                             vllm__module_executor__models__deepseek_v2__DeepseekV2DecoderLayer__init__)
MluHijackObject.apply_hijack(DeepseekV2ForCausalLM,
                             DeepseekV2ForCausalLM.load_weights,
                             vllm__module_executor__models__deepseek_v2__DeepseekV2ForCausalLM__load_weights)
