import math
import torch

from torch import nn
from typing import List, Optional, Union
from vllm.attention import Attention, AttentionMetadata
from vllm.config import CacheConfig
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from torch.nn import LayerNorm
from transformers import FalconConfig as HF_FalconConfig
from vllm.transformers_utils.configs import RWConfig
FalconConfig = Union[HF_FalconConfig, RWConfig]
from vllm_mlu.model_executor.layers.feed_forward import FeedForward
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.logger import init_logger
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.layers.linear import (QKVParallelLinear,
                                               RowParallelLinear)
from vllm.distributed import (get_tensor_model_parallel_rank,
                              get_tensor_model_parallel_world_size)

from vllm.model_executor.models.falcon import (FalconAttention,
                                               FalconDecoderLayer,
                                               _get_alibi_slopes)

logger = init_logger(__name__)


def vllm__module_executor__models__falcon__FalconAttention____init__(
    self,
    config: FalconConfig,
    cache_config: Optional[CacheConfig] = None,
    quant_config: Optional[QuantizationConfig] = None,
):
    super(FalconAttention, self).__init__()

    self.hidden_size = config.hidden_size
    tp_size = get_tensor_model_parallel_world_size()

    self.total_num_heads = config.num_attention_heads
    assert self.total_num_heads % tp_size == 0
    self.num_heads = self.total_num_heads // tp_size
    self.head_dim = self.hidden_size // self.total_num_heads
    assert self.head_dim * self.total_num_heads == self.hidden_size

    self.new_decoder_architecture = config.new_decoder_architecture
    self.multi_query = config.multi_query

    if self.new_decoder_architecture:
        self.total_num_kv_heads = config.num_kv_heads
    elif self.multi_query:
        self.total_num_kv_heads = 1
    else:
        self.total_num_kv_heads = self.total_num_heads
    if self.total_num_kv_heads >= tp_size:
        # Number of KV heads is greater than TP size, so we partition
        # the KV heads across multiple tensor parallel GPUs.
        assert self.total_num_kv_heads % tp_size == 0
    else:
        # Number of KV heads is less than TP size, so we replicate
        # the KV heads across multiple tensor parallel GPUs.
        assert tp_size % self.total_num_kv_heads == 0
    self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

    self.query_key_value = QKVParallelLinear(
        self.hidden_size,
        self.head_dim,
        self.total_num_heads,
        self.total_num_kv_heads,
        bias=config.bias,
        skip_bias_add=True,
        quant_config=quant_config,
    )
    self.q_size = self.num_heads * self.head_dim
    self.kv_size = self.num_kv_heads * self.head_dim

    # Layer-wise attention scaling
    self.inv_norm_factor = 1.0 / math.sqrt(self.head_dim)
    self.reduce_row_parallel_results = not (config.new_decoder_architecture
                                            or config.parallel_attn)
    self.dense = RowParallelLinear(
        self.hidden_size,
        self.hidden_size,
        bias=config.bias,
        skip_bias_add=True,
        quant_config=quant_config,
        reduce_results=self.reduce_row_parallel_results)

    self.use_rotary = config.rotary
    self.use_alibi = config.alibi
    assert not (self.use_rotary and self.use_alibi), (
        "Rotary and alibi are mutually exclusive.")
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: set cache_config for rotary & alibi
    ''' 
    if self.use_rotary:
        rope_theta = getattr(config, "rope_theta", 10000)
        max_position_embeddings = getattr(config,
                                            "max_position_embeddings", 8192)
        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=max_position_embeddings,
            base=rope_theta,
        )
        self.attn = Attention(self.num_heads,
                                self.head_dim,
                                self.inv_norm_factor,
                                num_kv_heads=self.num_kv_heads,
                                cache_config=cache_config,
                                quant_config=quant_config)
    elif self.use_alibi:
        tp_rank = get_tensor_model_parallel_rank()
        head_start = tp_rank * self.num_heads
        head_end = (tp_rank + 1) * self.num_heads
        alibi_slopes = (_get_alibi_slopes(self.total_num_heads) *
                        self.inv_norm_factor)
        alibi_slopes = alibi_slopes[head_start:head_end].tolist()
        self.attn = Attention(self.num_heads,
                                self.head_dim,
                                self.inv_norm_factor,
                                num_kv_heads=self.num_kv_heads,
                                alibi_slopes=alibi_slopes,
                                cache_config=cache_config,
                                quant_config=quant_config)
    else:
        self.attn = Attention(self.num_heads,
                                self.head_dim,
                                scale=self.inv_norm_factor,
                                num_kv_heads=self.num_kv_heads,
                                cache_config=cache_config,
                                quant_config=quant_config)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''


def vllm__module_executor__models__falcon__FalconAttention__forward(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
    kv_cache: torch.Tensor,
    attn_metadata: AttentionMetadata,
) -> torch.Tensor:
    qkv, bias = self.query_key_value(hidden_states)
    if bias is not None:
        qkv += bias
    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
    ''' 
    if self.use_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
    ==================
    '''
    attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
    attn_output, bias = self.dense(attn_output)
    return attn_output, bias


def vllm__module_executor__models__falcon__FalconDecoderLayer____init__(
    self,
    config: FalconConfig,
    cache_config: Optional[CacheConfig] = None,
    quant_config: Optional[QuantizationConfig] = None,
):
    super(FalconDecoderLayer, self).__init__()
    hidden_size = config.hidden_size
    self.num_heads = config.num_attention_heads
    self.self_attention = FalconAttention(config, cache_config,
                                            quant_config)
    '''
    =============================
    Modify by vllm_mlu
    =============================
    @brief: use FeedForward instead of MLP
    '''
    self.reduce_row_parallel_results = not (config.new_decoder_architecture
                                            or config.parallel_attn)
    self.mlp = FeedForward(hidden_size=hidden_size,
                            intermediate_size=hidden_size * 4,
                            hidden_act='gelu',
                            up_proj_name='dense_h_to_4h',
                            is_gated=False,
                            down_proj_name='dense_4h_to_h',
                            bias=config.bias,
                            quant_config=quant_config,
                            skip_bias_add=True,
                            reduce_results=self.reduce_row_parallel_results)
    '''
    ==================
    End of MLU Hijack
    ==================
    '''
    self.config = config

    if (config.num_ln_in_parallel_attn is None
            and config.new_decoder_architecture):
        config.num_ln_in_parallel_attn = 2

    if not config.parallel_attn:
        self.post_attention_layernorm = LayerNorm(
            hidden_size, eps=config.layer_norm_epsilon)
        self.input_layernorm = LayerNorm(hidden_size,
                                            eps=config.layer_norm_epsilon)
    else:
        if config.num_ln_in_parallel_attn == 2:
            # The layer norm before self-attention
            self.ln_attn = LayerNorm(hidden_size,
                                        eps=config.layer_norm_epsilon)
            # The layer norm before the MLP
            self.ln_mlp = LayerNorm(hidden_size,
                                    eps=config.layer_norm_epsilon)
        else:
            self.input_layernorm = LayerNorm(hidden_size,
                                                eps=config.layer_norm_epsilon)

    self.reduce_row_parallel_results = not (config.new_decoder_architecture
                                            or config.parallel_attn)


MluHijackObject.apply_hijack(FalconAttention,
                             FalconAttention.__init__,
                             vllm__module_executor__models__falcon__FalconAttention____init__)
MluHijackObject.apply_hijack(FalconAttention,
                             FalconAttention.forward,
                             vllm__module_executor__models__falcon__FalconAttention__forward)
MluHijackObject.apply_hijack(FalconDecoderLayer,
                             FalconDecoderLayer.__init__,
                             vllm__module_executor__models__falcon__FalconDecoderLayer____init__)
