import torch
import re
from typing import List, Optional, Tuple, Iterable, Union
from transformers import PretrainedConfig
from vllm.attention import AttentionMetadata
from vllm.config import CacheConfig
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.distributed import tensor_model_parallel_all_reduce
from vllm.model_executor.models.hunyuan import HunYuanAttention, HunYuanDecoderLayer, HunYuanForCausalLM, HunYuanModel
from vllm.model_executor.models.utils import is_pp_missing_parameter
from vllm.sequence import IntermediateTensors
from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
    get_compressed_tensors_cache_scale)
from vllm.model_executor.model_loader.weight_utils import (
    default_weight_loader, maybe_remap_kv_scale_name)
from vllm_mlu.model_executor.layers.sparse_moe_mlp import SparseMoeMlp
from vllm_mlu.model_executor.layers.feed_forward import FeedForward
from vllm_mlu.model_executor.models.layer_utils import (
    hunyuan_decoder_layer_forward_base, hunyuan_decoder_model_forward_base_pp,
    is_per_tensor_smoothquant, is_per_token_smoothquant, quant_fusion_with_rmsnorm)
from vllm_mlu.mlu_hijack_utils import MluHijackObject


class HunYuanSparseMoeBlock(SparseMoeMlp):

    def __init__(
        self,
        config: PretrainedConfig,
        quant_config: Optional[QuantizationConfig] = None,
    ):
        super().__init__(num_experts=config.num_experts,
                         top_k=config.moe_topk,
                         hidden_size=config.hidden_size,
                         intermediate_size=config.intermediate_size,
                         up_proj_name="gate_up_proj",
                         is_gated=True,
                         down_proj_name="down_proj",
                         has_bias=False,
                         skip_bias_add=False,
                         renormalize=True if config.moe_topk>1 else False,
                         hidden_act=config.hidden_act,
                         params_dtype=None,
                         quant_config=quant_config,
                         is_use_fused_moe=True)
        self.config = config
        self.shared_mlp = None
        if config.use_mixed_mlp_moe > 0:
            self.shared_mlp = FeedForward(hidden_size=config.hidden_size,
                                             intermediate_size=config.intermediate_size * config.num_shared_expert,
                                             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)


    def forward(self, hidden_states: torch.Tensor, residual: Optional[torch.Tensor] = None) -> torch.Tensor:
        num_tokens, hidden_dim = hidden_states.shape
        hidden_states = hidden_states.view(-1, hidden_dim)
        shared_output = None
        if self.shared_mlp is not None:
            shared_output = self.shared_mlp(hidden_states)

        # router_logits: (num_tokens, n_experts)
        router_logits, _ = self.gate(hidden_states)
        final_hidden_states = self.forward_experts(hidden_states, router_logits, residual)

        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 vllm__module_executor__models__hunyuan__HunYuanAttention__forward(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
    kv_cache: torch.Tensor,
    attn_metadata: AttentionMetadata,
    residual: Optional[torch.Tensor] = None,
    kv_states: Optional[Tuple[torch.Tensor]] = None,
    smooth_quant_scale: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    if self.attention_type == "self":
        qkv, _ = self.qkv_proj(hidden_states, smooth_quant_scale)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: pack q & k to fit tmo.apply_rotary
        '''
        qk, _ = qkv.split([self.q_size + self.kv_size, self.kv_size], dim=-1)
        self.rotary_emb(positions, qk.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
        '''
        ==================
        End of MLU Hijack
        ==================
        '''
        ori_k = k
        if self.use_qk_norm:
            q = self.query_layernorm(q.reshape(-1, self.num_heads, self.head_dim).contiguous()).reshape(-1, self.num_heads*self.head_dim)
            k = self.key_layernorm(k.reshape(-1, self.num_kv_heads, self.head_dim).contiguous()).reshape(-1, self.num_kv_heads*self.head_dim)
    elif self.attention_type == "cross":
        assert kv_states is not None
        ori_k, v = kv_states # use last layer kv,
        k = ori_k
        q, _ = self.q_proj(hidden_states, smooth_quant_scale)
        k_tmp = torch.empty_like(k) # Todo: reduant rotary embedding
        qk_temp = torch.cat((q, k_tmp), dim=-1)
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: pack q & k to fit tmo.apply_rotary
        '''
        self.rotary_emb(positions, qk_temp.view(-1, self.num_heads + self.num_kv_heads, self.head_dim))
        '''
        ==================
        End of MLU Hijack
        ==================
        '''
        if self.use_qk_norm:
            q = self.query_layernorm(q.view(-1, self.num_heads, self.head_dim).contiguous()).reshape(-1, self.num_heads*self.head_dim)
            k = self.key_layernorm(k.view(-1, self.num_kv_heads, self.head_dim).contiguous()).reshape(-1, self.num_kv_heads*self.head_dim)
    else:
        raise RuntimeError("Not support attnention type")
    attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: add residual
    '''
    output, _ = self.o_proj(attn_output, residual)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    return output, (ori_k, v)


def vllm__module_executor__models__hunyuan__HunYuanDecoderLayer____init__(
    self,
    config: PretrainedConfig,
    cache_config: Optional[CacheConfig] = None,
    quant_config: Optional[QuantizationConfig] = None,
    prefix: str = "",
    layer_id: int = -1,
) -> None:
    super(HunYuanDecoderLayer, self).__init__()
    self.hidden_size = config.hidden_size
    rope_theta = getattr(config, "rope_theta", 10000)
    rope_scaling = getattr(config, "rope_scaling", None)
    if rope_scaling is not None and getattr(
                config, "original_max_position_embeddings", None):
            rope_scaling["original_max_position_embeddings"] = (
                config.original_max_position_embeddings)
    max_position_embeddings = getattr(config, "max_position_embeddings",
                                        8192)
    # Support abacusai/Smaug-72B-v0.1 with attention_bias
    # Support internlm/internlm-7b with bias
    attention_bias = getattr(config, "attention_bias", False) or getattr(
        config, "bias", False)
    cla_factor = getattr(config, "cla_share_factor", 1)
    attention_type = "cross" \
        if layer_id >= 0 and layer_id % cla_factor != 0 else "self"
    self.self_attn = HunYuanAttention(
        config=config,
        hidden_size=self.hidden_size,
        num_heads=config.num_attention_heads,
        num_kv_heads=getattr(config, "num_key_value_heads",
                            config.num_attention_heads),
        rope_theta=rope_theta,
        rope_scaling=rope_scaling,
        max_position_embeddings=max_position_embeddings,
        quant_config=quant_config,
        bias=attention_bias,
        cache_config=cache_config,
        prefix=f"{prefix}.self_attn",
        attention_type=attention_type,
    )

    if getattr(config, "num_experts", None):
        self.mlp = HunYuanSparseMoeBlock(config=config,
                                        quant_config=quant_config)
    else:
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: use FeedForward instead of MLP
        '''
        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=getattr(config, "mlp_bias", False),
                               prefix=f"{prefix}.mlp",
                               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)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: prepare to perf per-tensor sq cases if suitable. For moe
        model, we only do quant fusion in attn block.
    '''
    self.is_per_tensor_sq_perf_cases = is_per_tensor_smoothquant(quant_config)
    self.is_per_token_sq_perf_cases = is_per_token_smoothquant(quant_config)
    if self.is_per_tensor_sq_perf_cases or self.is_per_token_sq_perf_cases:
        if self.self_attn.attention_type == "self":
            self.self_attn.qkv_proj.quant_method.skip_quant_input = True
        if self.self_attn.attention_type == "cross":
            self.self_attn.q_proj.quant_method.skip_quant_input = True
        self.quant_fusion_attn_layernorm = None
    '''
    ==================
    End of MLU Hijack
    ==================
    '''


def vllm__module_executor__models__hunyuan__HunYuanForCausalLM__load_weights(
        self,
        weights: Iterable[Tuple[str, torch.Tensor]]):
    cla_factor = getattr(self.config, "cla_share_factor", 1)
    '''
    =============================
    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)
        (".qkv_proj", ".q_proj", "q"),
        (".qkv_proj", ".k_proj", "k"),
        (".qkv_proj", ".v_proj", "v"),
        (".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
        if ("rotary_emb.cos_cached" in name
                or "rotary_emb.sin_cached" in name):
            # Models trained using ColossalAI may include these tensors in
            # the checkpoint. Skip them.
            continue
        # With tie_word_embeddings, we can skip lm_head.weight
        # The weight might appear unnecessarily in the files if the model is
        # processed with quantization, LoRA, fine-tuning, etc.
        if self.config.tie_word_embeddings and "lm_head.weight" in name:
            continue
        if scale_name := get_compressed_tensors_cache_scale(name):
            # Loading kv cache scales for compressed-tensors quantization
            param = params_dict[scale_name]
            weight_loader = getattr(param, "weight_loader",
                                    default_weight_loader)
            loaded_weight = loaded_weight[0]
            weight_loader(param, loaded_weight)
            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:
            if weight_name not in name:
                continue
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: delete if "mlp.experts" in name: continue condition
            '''
            '''
            ==================
            End of MLU Hijack
            ==================
            '''
            # cross layer only have q_proj, skip qkv pack
            if weight_name == ".q_proj":
                match = re.search(r'layers\.\d+', name)
                if match:
                    layer_id = int(match.group(0).split('.')[-1])
                    if cla_factor > 1 and layer_id % cla_factor != 0:
                        continue
            name = name.replace(weight_name, param_name)
            # 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
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: add expert skiped condition and delete useless if name not in params_dict: continue condition
            '''
            # Skip experts that are not assigned to this worker.
            if (("mlp.experts." in name or "mlp.shared_mlp." in name)
                    and name not in params_dict):
                continue
            '''
            ==================
            End of MLU Hijack
            ==================
            '''
            param = params_dict[name]
            weight_loader = param.weight_loader
            weight_loader(param, loaded_weight, shard_id)
            break
        else:
            if (name.endswith(".bias") and name not in params_dict):
                continue
            # Remapping the name of FP8 kv-scale.
            name = maybe_remap_kv_scale_name(name, params_dict)
            if name is None:
                continue

            if is_pp_missing_parameter(name, self):
                continue

            if "mlp.gate.wg." in name:
                name = name.replace("wg.", "")
            '''
            =============================
            Modify by vllm_mlu
            =============================
            @brief: add expert skiped condition
            '''
            # Skip experts that are not assigned to this worker.
            if (("mlp.experts." in name or "mlp.shared_mlp." in name)
                    and name not in params_dict):
                continue
            '''
            ==================
            End of MLU Hijack
            ==================
            '''
            param = params_dict[name]
            weight_loader = getattr(param, "weight_loader",
                                    default_weight_loader)
            weight_loader(param, loaded_weight)


def vllm__module_executor__models__hunyuan__HunYuanDecoderLayer__forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: AttentionMetadata,
        kv_states: Optional[Tuple[torch.Tensor]] = None,
        residual: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: perf model by:
    1) add residual in matmul;
    2) fuse quantization in layernorm in per-tensor sq case;
    '''
    attn_layernorm = self.input_layernorm
    if self.is_per_tensor_sq_perf_cases:
        if self.quant_fusion_attn_layernorm is None:
            if self.self_attn.attention_type == "self":
                self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
                    self.input_layernorm, self.self_attn.qkv_proj.scale_to_int)
            if self.self_attn.attention_type == "cross":
                self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
                    self.input_layernorm, self.self_attn.q_proj.scale_to_int)
        attn_layernorm = self.quant_fusion_attn_layernorm
    elif self.is_per_token_sq_perf_cases:
        if self.quant_fusion_attn_layernorm is None:
            if self.self_attn.attention_type == "self":
                self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
                    self.input_layernorm, self.self_attn.qkv_proj.smooth, dynamic_quant=True)
            if self.self_attn.attention_type == "cross":
                self.quant_fusion_attn_layernorm = quant_fusion_with_rmsnorm(
                    self.input_layernorm, self.self_attn.q_proj.smooth, dynamic_quant=True)
        attn_layernorm = self.quant_fusion_attn_layernorm
    return hunyuan_decoder_layer_forward_base(positions=positions,
                                      hidden_states=hidden_states,
                                      kv_cache=kv_cache,
                                      attn_metadata=attn_metadata,
                                      input_layernorm=attn_layernorm,
                                      self_attn=self.self_attn,
                                      post_layernorm=self.post_attention_layernorm,
                                      mlp=self.mlp,
                                      kv_states=kv_states,
                                      input_norm_fuse_en=self.is_per_token_sq_perf_cases)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''


def vllm__module_executor__models__hunyuan__HunYuanModel__forward(
        self,
        input_ids: Optional[torch.Tensor],
        positions: torch.Tensor,
        kv_caches: List[torch.Tensor],
        attn_metadata: AttentionMetadata,
        intermediate_tensors: Optional[IntermediateTensors],
        inputs_embeds: Optional[torch.Tensor] = None,
    ) -> Union[torch.Tensor, IntermediateTensors]:
    '''
    =============================
    Modify by vllm_mlu
    =============================
    '''
    return hunyuan_decoder_model_forward_base_pp(config=self.config,
                                         input_ids=input_ids,
                                         positions=positions,
                                         kv_caches=kv_caches,
                                         attn_metadata=attn_metadata,
                                         intermediate_tensors=intermediate_tensors,
                                         layers=self.layers,
                                         start_layer=self.start_layer,
                                         end_layer=self.end_layer,
                                         get_input_embeddings=self.get_input_embeddings,
                                         norm=self.norm,
                                         inputs_embeds=inputs_embeds)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''


MluHijackObject.apply_hijack(HunYuanAttention,
                             HunYuanAttention.forward,
                             vllm__module_executor__models__hunyuan__HunYuanAttention__forward)
MluHijackObject.apply_hijack(HunYuanDecoderLayer,
                             HunYuanDecoderLayer.__init__,
                             vllm__module_executor__models__hunyuan__HunYuanDecoderLayer____init__)
MluHijackObject.apply_hijack(HunYuanForCausalLM,
                             HunYuanForCausalLM.load_weights,
                             vllm__module_executor__models__hunyuan__HunYuanForCausalLM__load_weights)
MluHijackObject.apply_hijack(HunYuanDecoderLayer,
                             HunYuanDecoderLayer.forward,
                             vllm__module_executor__models__hunyuan__HunYuanDecoderLayer__forward)
MluHijackObject.apply_hijack(HunYuanModel,
                             HunYuanModel.forward,
                             vllm__module_executor__models__hunyuan__HunYuanModel__forward)
