# SPDX-License-Identifier: Apache-2.0
# Derived from vLLM; modified by DeepSeek.
# Adapted from
# https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/llama/modeling_llama.py
# Copyright 2023 The vLLM team.
# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
#
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
# and OPT implementations in this library. It has been modified from its
# original forms to accommodate minor architectural differences compared
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import math
from typing import Any

import torch
import torch.nn as nn
from tile_kernels.transform import apply_rotary


def _build_cos_sin_cache(t: torch.Tensor, inv_freq: torch.Tensor, mscale: float = 1.0) -> torch.Tensor:
    # Preserve the reference implementation's order: polar conversion before scaling.
    freqs = torch.outer(t, inv_freq)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    cos = freqs_cis.real.cuda() * mscale
    sin = freqs_cis.imag.cuda() * mscale
    return torch.cat((cos, sin), dim=-1)


class RotaryEmbedding(nn.Module):
    """Original rotary positional embedding."""

    mscale: float = 1.0
    """cos/sin 幅值缩放，YaRN 子类在 __init__ 中覆盖"""

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

        cache = self._compute_cos_sin_cache()
        assert cache.dtype == torch.float
        self.register_buffer("cos_sin_cache", cache, persistent=False)

    def _compute_inv_freq(self) -> torch.Tensor:
        """Compute the inverse frequency."""
        # NOTE(woosuk): The HF implementation uses `torch.arange(...).float()`.
        # However, we use `torch.arange(..., dtype=torch.float)` instead to
        # avoid numerical issues with large base values (e.g., 10000000).
        # This may cause a slight numerical difference between the HF
        # implementation and ours.
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.rotary_dim, 2, dtype=torch.float, device="cuda") / self.rotary_dim))
        return inv_freq

    def _compute_cos_sin_cache(self) -> torch.Tensor:
        """Compute the cos and sin cache."""
        inv_freq = self._compute_inv_freq()
        t = torch.arange(self.max_position_embeddings, dtype=torch.float, device="cuda")
        return _build_cos_sin_cache(t, inv_freq, mscale=self.mscale)

    def forward(
        self,
        positions: torch.Tensor,
        query: torch.Tensor,
        key: torch.Tensor | None = None,
        conjugate: bool = False,
    ) -> None:
        assert positions.shape[0] == query.shape[0], f"{positions.shape=}, {query.shape=}"
        # `query` shape: [num_tokens, num_heads, qk_rope_head_dim]
        # `key` shape: [num_tokens, 1, qk_rope_head_dim]
        # `apply_rotary()` is an in-place operation that updates the query and key tensors.
        assert query.shape[-1] == self.rotary_dim
        if key is not None:
            assert key.shape[-1] == self.rotary_dim
        # TODO(NPU): When head_num = 1, positions = None is faster.
        apply_rotary(
            query,
            self.cos_sin_cache,
            key,
            positions=positions,
            interleaved=not self.is_neox_style,
            conjugate=conjugate,
        )


# Inverse dim formula to find dim based on number of rotations
def _yarn_find_correction_dim(
    num_rotations: int,
    dim: int,
    base: float = 10000,
    max_position_embeddings: int = 2048,
) -> float:
    return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (2 * math.log(base))


# Find dim range bounds based on rotations
def _yarn_find_correction_range(
    low_rot: int,
    high_rot: int,
    dim: int,
    base: float = 10000,
    max_position_embeddings: int = 2048,
) -> tuple[int, int]:
    low = math.floor(_yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings))
    high = math.ceil(_yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings))
    return max(low, 0), min(high, dim - 1)  # Clamp values just in case


def _yarn_linear_ramp_mask(
    low: float,
    high: float,
    dim: int,
    dtype: torch.dtype,
    device: torch.device,
) -> torch.Tensor:
    if low == high:
        high += 0.001  # Prevent singularity

    linear_func = (torch.arange(dim, dtype=dtype, device=device) - low) / (high - low)
    ramp_func = torch.clamp(linear_func, 0, 1)
    return ramp_func


def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
    if scale <= 1:
        return 1.0
    return 0.1 * mscale * math.log(scale) + 1.0


class YaRNScalingRotaryEmbedding(RotaryEmbedding):
    """RotaryEmbedding extended with YaRN method.

    Credits to Peng et al. github.com/jquesnelle/yarn
    """

    def __init__(
        self,
        head_size: int,
        rotary_dim: int,
        original_max_position_embeddings: int,
        max_position_embeddings: int,
        base: int,
        is_neox_style: bool,
        scaling_factor: float,
        *,
        extrapolation_factor: float = 1,
        attn_factor: float = 1,
        beta_fast: float = 32,
        beta_slow: float = 1,
        mscale: float | bool = True,
        mscale_all_dim: float = 0,
        rope_freq_sample_interval: int | None = None,
        high_freq_cycle_threshold: int | float | None = None,
    ) -> None:
        self.rope_freq_sample_interval = rope_freq_sample_interval  # inv_freq[::self.rope_freq_sample_interval]
        self.high_freq_cycle_threshold = high_freq_cycle_threshold  # inv_freq[high_freq_start_idx:]
        if self.high_freq_cycle_threshold is not None:
            self.inv_freq_hi = 2 * math.pi / self.high_freq_cycle_threshold

        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
        )
        self.original_max_position_embeddings = original_max_position_embeddings
        super().__init__(head_size, rotary_dim, max_position_embeddings, base, is_neox_style)

    def _compute_inv_freq(self) -> torch.Tensor:
        pos_freqs = self.base ** (torch.arange(0, self.rotary_dim, 2, dtype=torch.float, device="cuda") / self.rotary_dim)
        inv_freq_extrapolation = 1.0 / pos_freqs
        inv_freq_interpolation = 1.0 / (self.scaling_factor * pos_freqs)

        low, high = _yarn_find_correction_range(
            self.beta_fast, self.beta_slow, self.rotary_dim, self.base, self.original_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, device="cuda")
        ) * self.extrapolation_factor
        inv_freq = inv_freq_interpolation * (1 - inv_freq_mask) + inv_freq_extrapolation * inv_freq_mask

        # === sample inv_freq ===
        inv_freq_idx_s = 0
        inv_freq_sample_interval = self.rope_freq_sample_interval or 1
        if self.high_freq_cycle_threshold is not None:
            assert inv_freq.equal(inv_freq.sort(descending=True).values)  # 降序
            inv_freq_idx_s = torch.where(inv_freq <= self.inv_freq_hi)[0][0]
        inv_freq = inv_freq[inv_freq_idx_s::inv_freq_sample_interval]
        # === end sample inv_freq ===

        return inv_freq


_ROPE_DICT: dict[tuple, RotaryEmbedding] = {}


def get_rope(
    head_size: int,
    rotary_dim: int,
    max_position: int,
    base: int,
    is_neox_style: bool = True,
    rope_scaling: dict[str, Any] | None = None,
) -> RotaryEmbedding:
    key = (
        head_size,
        rotary_dim,
        max_position,
        base,
        is_neox_style,
        tuple(rope_scaling.items()) if rope_scaling is not None else None,
    )
    if key in _ROPE_DICT:
        return _ROPE_DICT[key]

    if rope_scaling is None:
        rotary_emb = RotaryEmbedding(
            head_size,
            rotary_dim,
            max_position,
            base,
            is_neox_style,
        )
    else:
        scaling_type = rope_scaling["type"]
        if 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",
                    "mscale",
                    "mscale_all_dim",
                    "rope_freq_sample_interval",
                    "high_freq_cycle_threshold",
                )
            }
            rotary_emb = YaRNScalingRotaryEmbedding(
                head_size,
                rotary_dim,
                original_max_position,
                max_position,
                base,
                is_neox_style,
                scaling_factor,
                **extra_kwargs,
            )
        else:
            raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
    _ROPE_DICT[key] = rotary_emb
    return rotary_emb