from typing import Any, Dict, List, Optional, Tuple, Union

import math
import torch

from vllm.attention import AttentionMetadata
from vllm.model_executor.custom_op import CustomOp
from vllm.model_executor.layers.rotary_embedding import (
    RotaryEmbedding, MRotaryEmbedding,
    LinearScalingRotaryEmbedding, DeepseekScalingRotaryEmbedding,
    DynamicNTKScalingRotaryEmbedding, DynamicNTKAlphaRotaryEmbedding,
    YaRNScalingRotaryEmbedding, Phi3LongRoPEScaledRotaryEmbedding,
    _yarn_find_correction_range, _ROPE_DICT, yarn_get_mscale, _yarn_linear_ramp_mask)
from vllm.model_executor.layers import rotary_embedding
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.envs import VLLM_ALLOW_LONG_MAX_MODEL_LEN
from vllm.logger import init_logger

logger = init_logger(__name__)

def get_long_max_model_max_position_emb(max_position_embeddings, scaling_factor):
    if MLURotaryEmbedding.max_seq_len != None and \
            MLURotaryEmbedding.max_seq_len > max_position_embeddings * scaling_factor:
        logger.warning(f"User-specified max_model_len ({MLURotaryEmbedding.max_seq_len}) is different with " +
                        f"max_position_embedding ({max_position_embeddings}) * scaling_factor ({scaling_factor}) " +
                        "from model's config.json, This may lead to incorrect model outputs or MLU errors. " +
                        f"Make sure the value is correct and within the model context size. " +
                        f"Set max_position_embedding={MLURotaryEmbedding.max_seq_len}.")
        return math.ceil(MLURotaryEmbedding.max_seq_len / scaling_factor)
    return max_position_embeddings


@CustomOp.register("rotary_embedding_mlu")
class MLURotaryEmbedding(RotaryEmbedding, CustomOp):

    cu_seq_lens : torch.Tensor = None
    max_seq_len : int = None
    is_prompt : bool = False
    is_chunked : bool = False
    set_cos_sin : bool = False
    cos_ : torch.Tensor = None
    sin_ : torch.Tensor = None
    positions_: torch.Tensor = None

    def __init__(
        self,
        head_size: int,
        rotary_dim: int,
        max_position_embeddings: int,
        base: int,
        is_neox_style: bool,
        dtype: torch.dtype,
    ) -> None:
        CustomOp.__init__(self)
        self.head_size = head_size
        self.rotary_dim = rotary_dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        self.is_neox_style = is_neox_style
        self.dtype = dtype

        if MLURotaryEmbedding.max_seq_len != None \
                and self.max_position_embeddings < MLURotaryEmbedding.max_seq_len and \
                not isinstance(self, (YaRNScalingRotaryEmbedding, DeepseekScalingRotaryEmbedding)):
            logger.warning(f"User-specified max_model_len ({MLURotaryEmbedding.max_seq_len}) is different with " +
                            f"max_position_embedding ({max_position_embeddings}) from model's config.json, " +
                            f"This may lead to incorrect model outputs or MLU errors. " +
                            f"Make sure the value is correct and within the model context size. " +
                            f"Set max_position_embedding={MLURotaryEmbedding.max_seq_len}.")
            self.max_position_embeddings = MLURotaryEmbedding.max_seq_len
        cache = self._compute_cos_sin_cache()
        if isinstance(self, MLULinearScalingRotaryEmbedding):
            logger.debug(f"Using mlu defining _compute_cos_sin_cache due to the special tensor composition")
        elif is_neox_style:
            cache_pos = cache.shape[0]
            cache = cache.reshape(cache_pos, 2, -1)
            cache = torch.tile(cache, (1, 1, 2)).reshape(cache_pos, -1)
        else:
            cache = cache.repeat_interleave(2, dim=-1)

        cache = cache.to(dtype)
        self.cos_sin_cache: torch.Tensor
        self.register_buffer("cos_sin_cache", cache, persistent=False)


    @classmethod
    def set_mlu_var(
        cls,
        input_ids: torch.Tensor,
        attn_metadata: AttentionMetadata
    ) -> None:
        cls.unset_mlu_var()
        is_chunked = False
        is_prompt = False
        prefill_metadata = attn_metadata.prefill_metadata
        decode_metadata = attn_metadata.decode_metadata
        if prefill_metadata:
            cu_seq_lens = prefill_metadata.query_start_loc
            rope_max_seq_len = prefill_metadata.max_query_len
            is_prompt = True
            # Workaround: mlugraph does not support torch.ne|eq|equal .etc for now,
            # because context mlugraph always uses in benchmark latency, and in this
            # case, query_start_loc always equals to seq_start_loc, so we can set
            # is_chunked to False directly.
            if prefill_metadata.use_cuda_graph:
                is_chunked = False
            elif decode_metadata or \
                max(prefill_metadata.seq_lens) != prefill_metadata.max_query_len:
                is_chunked = True
        if decode_metadata:
            if prefill_metadata:
                cu_seq_lens = attn_metadata.query_start_loc
                rope_max_seq_len = max(rope_max_seq_len,
                    attn_metadata.num_prefill_tokens + attn_metadata.num_decode_tokens)
            else:
                # input_ids is pack mode, and the decode_seq_len = 1
                cu_seq_lens = torch.arange(0, input_ids.shape[0] + 1, 1, dtype=torch.int32, device="mlu")
                rope_max_seq_len = 1
        cls.cu_seq_lens = cu_seq_lens
        cls.max_seq_len = rope_max_seq_len
        cls.is_prompt = is_prompt
        cls.is_chunked = is_chunked

    @classmethod
    def unset_mlu_var(cls):
        cls.cu_seq_lens = None
        cls.max_seq_len = None
        cls.is_prompt = False
        cls.is_chunked = False
        cls.set_cos_sin = False
        cls.cos_ = None
        cls.sin_ = None
        cls.positions_ = None

    def _get_cos_sin(self) -> Tuple[torch.Tensor, torch.Tensor]:
        cos, sin = self.cos_sin_cache.chunk(2, dim=-1)
        sin = sin.view(-1, self.rotary_dim)
        cos = cos.view(-1, self.rotary_dim)
        return cos, sin

    def _get_positions_with_offsets_mlu(
        self,
        positions: torch.Tensor,
        offsets: torch.Tensor
    ) -> torch.Tensor:
        if offsets.numel() != positions.numel():
            raise Exception("rope offsets numel mismatch with positions, "
                            f"positions: {positions.numel()}, offsets: {offsets.numel()}")
        return (positions + offsets).to(torch.int32)

    def forward_mlu(
        self,
        positions: torch.Tensor,
        x: torch.Tensor,
        offsets: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        from vllm import _mlu_ops as mlu_ops

        # ops.rotary_embedding()/batched_rotary_embedding()
        # are in-place operations that update the query and key tensors.
        if MLURotaryEmbedding.set_cos_sin == False:
            MLURotaryEmbedding.cos_, MLURotaryEmbedding.sin_ = self._get_cos_sin()
            MLURotaryEmbedding.set_cos_sin = True
        interleaved = True
        if self.is_neox_style:
            interleaved = False

        if offsets is not None:
            if MLURotaryEmbedding.positions_ is None:
                MLURotaryEmbedding.positions_ = (
                    self._get_positions_with_offsets_mlu(positions, offsets))
            position_ids = MLURotaryEmbedding.positions_
            discrete = True
        elif MLURotaryEmbedding.is_chunked or not MLURotaryEmbedding.is_prompt:
            position_ids = positions
            discrete = True
        else:
            position_ids = None
            discrete = False

        x = mlu_ops.rotary_embedding(x,
                                     MLURotaryEmbedding.sin_,
                                     MLURotaryEmbedding.cos_,
                                     position_ids,
                                     MLURotaryEmbedding.cu_seq_lens,
                                     interleaved,
                                     discrete,
                                     False,
                                     MLURotaryEmbedding.max_seq_len)
        return x


class MLULinearScalingRotaryEmbedding(MLURotaryEmbedding, LinearScalingRotaryEmbedding):

    def __init__(
        self,
        head_size: int,
        rotary_dim: int,
        max_position_embeddings: int,
        base: int,
        is_neox_style: bool,
        scaling_factors: Union[List[float], float],
        dtype: torch.dtype,
    ) -> None:
        if isinstance(scaling_factors, float):
            scaling_factors = [scaling_factors]
        self.scaling_factors: List[float] = scaling_factors  # noqa
        MLURotaryEmbedding.__init__(self, head_size, rotary_dim,
                                    max_position_embeddings, base,
                                    is_neox_style, dtype)
        # Lazy initialized.
        self._scaling_factor_to_offset: Dict[float, int]

    def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:
        """Compute the inverse frequency."""
        half_dim = self.rotary_dim // 2
        if self.is_neox_style:
            inv_freq = 1.0 / (base ** ((torch.arange(
                0, self.rotary_dim, 1, dtype=torch.float32, device="mlu") % half_dim) * 2 / self.rotary_dim)
            )
        else:
            inv_freq = 1.0 / (
                    base
                    ** ( torch.arange(0, self.rotary_dim, 1, device="mlu", dtype=torch.float32) // 2 * 2
                        / self.rotary_dim
                    )
            )
        return inv_freq

    def _compute_cos_sin_cache(self) -> torch.Tensor:
        inv_freq = self._compute_inv_freq(self.base)
        cache_list: List[torch.Tensor] = []
        # offsets to the next cache in a tensor.
        # Each offset corresponds to the same index in scaling_factors.
        offsets: List[int] = []
        for scaling_factor in self.scaling_factors:
            # NOTE(woosuk): self.max_position_embeddings is the original
            # maximum length before applying the rope scaling.
            # Thus, the maximum length after applying the rope scaling is
            # self.max_position_embeddings * self.scaling_factor.
            max_len = self.max_position_embeddings * scaling_factor
            t = torch.arange(max_len, dtype=torch.float, device="mlu")
            t = t / scaling_factor

            freqs = torch.einsum("i,j -> ij", t, inv_freq)
            cos = freqs.cos()
            sin = freqs.sin()
            cache = torch.cat((cos, sin), dim=-1)
            if not cache_list:
                offset = 0
            else:
                last_offset = offsets[-1]
                next_max_len = cache_list[-1].shape[0]
                offset = last_offset + next_max_len
            offsets.append(offset)
            cache_list.append(cache)
        self._scaling_factor_to_offset = {
            float(scaling_factor): offsets[i]
            for i, scaling_factor in enumerate(self.scaling_factors)
        }
        assert len(self.scaling_factors) == len(offsets)
        return torch.cat(cache_list, dim=0)


class MLUDeepseekScalingRotaryEmbedding(MLURotaryEmbedding, DeepseekScalingRotaryEmbedding):

    def __init__(
        self,
        head_size: int,
        rotary_dim: int,
        max_position_embeddings: int,
        base: int,
        is_neox_style: bool,
        scaling_factor: float,
        dtype: torch.dtype,
        *,
        extrapolation_factor: float = 1,
        attn_factor: float = 1,
        beta_fast: int = 32,
        beta_slow: int = 1,
        mscale: float = 1,
        mscale_all_dim: float = 0,
    ) -> None:
        self.scaling_factor = scaling_factor
        self.extrapolation_factor = extrapolation_factor
        self.attn_factor = attn_factor
        self.beta_fast = beta_fast
        self.beta_slow = beta_slow
        # Get n-d magnitude scaling corrected for interpolation.
        self.mscale = float(
            yarn_get_mscale(self.scaling_factor, float(mscale)) /
            yarn_get_mscale(self.scaling_factor, float(mscale_all_dim)) *
            attn_factor)
        MLURotaryEmbedding.__init__(self, head_size, rotary_dim,
                                    max_position_embeddings, base,
                                    is_neox_style, dtype)

    def forward_mlu(
        self,
        positions: torch.Tensor,
        query: torch.Tensor,
        key: torch.Tensor,
        offsets: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        from vllm import _mlu_ops as mlu_ops

        """PyTorch-native implementation equivalent to forward()."""
        query_rot = query[..., :self.rotary_dim]
        key_rot = key[..., :self.rotary_dim]
        if self.rotary_dim < self.head_size:
            query_pass = query[..., self.rotary_dim:]
            key_pass = key[..., self.rotary_dim:]
        if MLURotaryEmbedding.set_cos_sin == False:
            MLURotaryEmbedding.cos_, MLURotaryEmbedding.sin_ = self._get_cos_sin()
            MLURotaryEmbedding.set_cos_sin = True
        interleaved = True
        if self.is_neox_style:
            interleaved = False
        if MLURotaryEmbedding.is_chunked or not MLURotaryEmbedding.is_prompt:
            position_ids = positions
            discrete = True
        else :
            position_ids = None
            discrete = False
        query_rot = mlu_ops.rotary_embedding(query_rot,
                                    MLURotaryEmbedding.sin_,
                                    MLURotaryEmbedding.cos_,
                                    position_ids,
                                    MLURotaryEmbedding.cu_seq_lens,
                                    interleaved,
                                    discrete,
                                    False,
                                    MLURotaryEmbedding.max_seq_len)
        key_rot = mlu_ops.rotary_embedding(key_rot,
                                    MLURotaryEmbedding.sin_,
                                    MLURotaryEmbedding.cos_,
                                    position_ids,
                                    MLURotaryEmbedding.cu_seq_lens,
                                    interleaved,
                                    discrete,
                                    False,
                                    MLURotaryEmbedding.max_seq_len)
        if self.rotary_dim < self.head_size:
            query = torch.cat((query_rot, query_pass), dim=-1)
            key = torch.cat((key_rot, key_pass), dim=-1)
        else:
            query = query_rot
            key = key_rot
        return query, key

    def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor:
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: change device cuda to mlu
        '''  
        pos_freqs = self.base**(torch.arange(
            0, self.rotary_dim, 2, dtype=torch.float, device="mlu") /
                                self.rotary_dim)
        '''
        ==================
        End of MLU Hijack
        ==================
        '''
        inv_freq_extrapolation = 1.0 / pos_freqs
        inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs)
    
        low, high = _yarn_find_correction_range(self.beta_fast, self.beta_slow,
                                                self.rotary_dim, self.base,
                                                self.max_position_embeddings)
        # Get n-d rotational scaling corrected for extrapolation
        inv_freq_mask = (1 - _yarn_linear_ramp_mask(
            low, high, self.rotary_dim // 2,
            dtype=torch.float)) * self.extrapolation_factor
        inv_freq = inv_freq_interpolation * (
            1 - inv_freq_mask) + inv_freq_extrapolation * inv_freq_mask
        return inv_freq
    
    def _compute_cos_sin_cache(self) -> torch.Tensor:
        inv_freq = self._compute_inv_freq(self.scaling_factor)
        '''
        =============================
        Modify by vllm_mlu
        =============================
        @brief: change device cuda to mlu
        '''  
        t = torch.arange(self.max_position_embeddings * self.scaling_factor,
                         device="mlu",
                         dtype=torch.float32)
        '''
        ==================
        End of MLU Hijack
        ==================
        '''
        freqs = torch.einsum("i,j -> ij", t, inv_freq)
        cos = (freqs.cos() * self.mscale)
        sin = (freqs.sin() * self.mscale)
        cache = torch.cat((cos, sin), dim=-1)
        print("Cache shape", cache.shape)
        return cache


class MLULlama3RotaryEmbedding(MLURotaryEmbedding):

    def __init__(
        self,
        head_size: int,
        rotary_dim: int,
        max_position_embeddings: int,
        base: int,
        is_neox_style: bool,
        dtype: torch.dtype,
        scaling_factor: float,
        low_freq_factor: float,
        high_freq_factor: float,
        orig_max_position: int,
    ) -> None:
        self.scaling_factor = scaling_factor
        self.low_freq_factor = low_freq_factor
        self.high_freq_factor = high_freq_factor
        self.orig_max_position = orig_max_position
        super().__init__(head_size, rotary_dim, max_position_embeddings, base,
                         is_neox_style, dtype)


class MLUDynamicNTKAlphaRotaryEmbedding(MLURotaryEmbedding, DynamicNTKAlphaRotaryEmbedding):
    """RotaryEmbedding extended with Dynamic NTK scaling.
    Credits to the Reddit users /u/bloc97 and /u/emozilla
    """

    def __init__(
        self,
        head_size: int,
        rotary_dim: int,
        max_position_embeddings: int,
        base: int,
        is_neox_style: bool,
        scaling_alpha: float,
        dtype: torch.dtype,
    ) -> None:
        self.scaling_alpha = scaling_alpha
        MLURotaryEmbedding.__init__(
            self, head_size, rotary_dim, max_position_embeddings, base,
            is_neox_style, dtype)


class MLUMRotaryEmbedding(MLURotaryEmbedding, MRotaryEmbedding):
    """Rotary Embedding with Multimodal Sections."""

    def __init__(
        self,
        head_size: int,
        rotary_dim: int,
        max_position_embeddings: int,
        base: int,
        is_neox_style: bool,
        dtype: torch.dtype,
        mrope_section: Optional[List[int]] = None,
    ) -> None:
        MLURotaryEmbedding.__init__(
            self, head_size, rotary_dim, max_position_embeddings, base,
            is_neox_style, dtype)

        self.mrope_section = mrope_section
        if self.mrope_section:
            assert sum(self.mrope_section) == rotary_dim // 2

    def forward_mlu(
        self,
        positions: torch.Tensor,
        x: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        assert positions.ndim == 1 or positions.ndim == 2
        num_tokens = positions.shape[-1]
        
        cos_sin = self.cos_sin_cache[positions]
        cos, sin = cos_sin.chunk(2, dim=-1)
        if positions.ndim == 2:
            assert self.mrope_section
            num_section = len(self.mrope_section)
            mrope_section = self.mrope_section * 2
            cos = torch.cat([
                m[i % num_section]
                for i, m in enumerate(cos.split(mrope_section, dim=-1))
            ],
                            dim=-1)
            sin = torch.cat([
                m[i % num_section]
                for i, m in enumerate(sin.split(mrope_section, dim=-1))
            ],
                            dim=-1)
        from vllm import _mlu_ops as mlu_ops
        interleaved = True
        if self.is_neox_style:
            interleaved = False
        position_ids = None
        discrete = False
        # mlu_ops.rotary_embedding() is a in-place operation that update the query and key tensors.
        x = mlu_ops.rotary_embedding(x,
                                     sin,
                                     cos,
                                     position_ids,
                                     MLURotaryEmbedding.cu_seq_lens,
                                     interleaved,
                                     discrete,
                                     False,
                                     MLURotaryEmbedding.max_seq_len)
        return x


def vllm__model_executor__layers__rotary_embedding__get_rope(
    head_size: int,
    rotary_dim: int,
    max_position: int,
    base: int,
    is_neox_style: bool = True,
    rope_scaling: Optional[Dict[str, Any]] = None,
    dtype: Optional[torch.dtype] = None,
    partial_rotary_factor: float = 1.0,
) -> RotaryEmbedding:
    if dtype is None:
        dtype = torch.get_default_dtype()
    if rope_scaling is not None:
        # Transforms every value that is a list into a tuple for caching calls
        rope_scaling_tuple = {
            k: tuple(v) if isinstance(v, list) else v
            for k, v in rope_scaling.items()
        }
        rope_scaling_args = tuple(rope_scaling_tuple.items())
    else:
        rope_scaling_args = None
    if partial_rotary_factor < 1.0:
        rotary_dim = int(rotary_dim * partial_rotary_factor)
    key = (head_size, rotary_dim, max_position, base, is_neox_style,
           rope_scaling_args, dtype)
    if key in _ROPE_DICT:
        return _ROPE_DICT[key]

    if rope_scaling is None:
        rotary_emb = MLURotaryEmbedding(head_size, rotary_dim, max_position, base,
                                        is_neox_style, dtype)
    else:
        scaling_type = rope_scaling["rope_type"]

        if scaling_type == "llama3":
            scaling_factor = rope_scaling["factor"]
            low_freq_factor = rope_scaling["low_freq_factor"]
            high_freq_factor = rope_scaling["high_freq_factor"]
            original_max_position = rope_scaling[
                "original_max_position_embeddings"]
            rotary_emb = MLULlama3RotaryEmbedding(head_size, rotary_dim,
                                                  max_position, base,
                                                  is_neox_style, dtype,
                                                  scaling_factor, low_freq_factor,
                                                  high_freq_factor,
                                                  original_max_position)
        elif scaling_type == "default":
            if "mrope_section" in rope_scaling:
                rotary_emb = MLUMRotaryEmbedding(
                    head_size,
                    rotary_dim,
                    max_position,
                    base,
                    is_neox_style,
                    dtype,
                    mrope_section=rope_scaling["mrope_section"],
                )
            else:
                rotary_emb = MLURotaryEmbedding(
                    head_size,
                    rotary_dim,
                    max_position,
                    base,
                    is_neox_style,
                    dtype,
                )
        elif scaling_type == "linear":
            scaling_factor = rope_scaling["factor"]
            rotary_emb = MLULinearScalingRotaryEmbedding(head_size, rotary_dim,
                                                         max_position, base,
                                                         is_neox_style,
                                                         scaling_factor, dtype)
        elif scaling_type == "dynamic":
            if "alpha" in rope_scaling:
                rotary_emb = MLUDynamicNTKAlphaRotaryEmbedding(
                    head_size, rotary_dim, max_position, base, is_neox_style,
                    rope_scaling["alpha"], dtype)
            else:
                scaling_factor = rope_scaling["factor"]
                rotary_emb = DynamicNTKScalingRotaryEmbedding(
                    head_size, rotary_dim, max_position, base, is_neox_style,
                    scaling_factor, dtype)
        elif scaling_type == "yarn":
            scaling_factor = rope_scaling["factor"]
            original_max_position = rope_scaling[
                "original_max_position_embeddings"]
            extra_kwargs = {
                k: v
                for k, v in rope_scaling.items()
                if k in ("extrapolation_factor", "attn_factor", "beta_fast",
                         "beta_slow")
            }
            original_max_position = get_long_max_model_max_position_emb(original_max_position, scaling_factor)
            rotary_emb = YaRNScalingRotaryEmbedding(head_size, rotary_dim,
                                                    original_max_position,
                                                    base, is_neox_style,
                                                    scaling_factor, dtype,
                                                    **extra_kwargs)
        elif scaling_type == "deepseek_yarn":
            scaling_factor = rope_scaling["factor"]
            original_max_position = rope_scaling[
                "original_max_position_embeddings"]
            # assert max_position == original_max_position * scaling_factor
            extra_kwargs = {
                k: v
                for k, v in rope_scaling.items()
                if k in ("extrapolation_factor", "attn_factor", "beta_fast",
                         "beta_slow", "mscale", "mscale_all_dim")
            }
            original_max_position = get_long_max_model_max_position_emb(original_max_position, scaling_factor)
            rotary_emb = MLUDeepseekScalingRotaryEmbedding(
                head_size, rotary_dim, original_max_position, base,
                is_neox_style, scaling_factor, dtype, **extra_kwargs)
        elif scaling_type == "longrope":
            short_factor = rope_scaling["short_factor"]
            long_factor = rope_scaling["long_factor"]
            original_max_position = rope_scaling[
                "original_max_position_embeddings"]
            extra_kwargs = {
                k: v
                for k, v in rope_scaling.items()
                if k in ("short_mscale", "long_mscale")
            }
            rotary_emb = Phi3LongRoPEScaledRotaryEmbedding(
                head_size, rotary_dim, max_position, original_max_position,
                base, is_neox_style, dtype, short_factor, long_factor,
                **extra_kwargs)
        else:
            raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
    _ROPE_DICT[key] = rotary_emb
    return rotary_emb


MluHijackObject.apply_hijack(rotary_embedding,
                             rotary_embedding.get_rope,
                             vllm__model_executor__layers__rotary_embedding__get_rope)
