import math
import torch
from typing import Optional, Tuple

from .model_mapper import ModelMapper
from .custom_op import FusedAttention, MoE

class Embedding(torch.nn.Module):
    def __init__(self, embed, config):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.embed = embed
        self.embed_scale = 1.0
        if config.model_type == 'gemma' or config.model_type == 'gemma2':
            self.embed_scale = self.hidden_size**0.5
        if hasattr(embed, 'embed_scale'):
            self.embed_scale = embed.embed_scale

    def forward(self, input_ids):
        inputs_embeds = self.embed(input_ids).view(-1, 1, self.hidden_size)
        return inputs_embeds

def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)

class Attention(torch.nn.Module):
    def __init__(self, attn, layer_id, config):
        super().__init__()
        self.export_fused_attn = False
        if config is None: return
        self.config = config
        self.fused_attn = FusedAttention(config.hidden_size, f'/layers.{layer_id}/self_attn/FusedAttention')
        self.layer_id = layer_id
        self.hidden_size = config.hidden_size
        self.head_dim = config.head_dim
        if isinstance(config.num_attention_heads, list):
            self.num_heads = config.num_attention_heads[layer_id]
            self.num_key_value_heads = config.num_key_value_heads[layer_id]
        else:
            self.head_dim = config.head_dim
            self.num_heads = config.num_attention_heads
            self.num_key_value_heads = config.num_key_value_heads
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.rotary = config.rotary

        ModelMapper.do_map(self, attn, config.model_map['attention'])

        if hasattr(self, 'qkv_proj') and self.qkv_proj is not None:
            # split qkv linear to q, k, v
            split_sizes = [self.hidden_size] * 3
            if self.qkv_proj.weight.shape[0] != self.hidden_size * 3:
                # M/GQA
                split_sizes = [
                    self.num_heads * self.head_dim,           # q_size
                    self.num_key_value_heads * self.head_dim, # k_size
                    self.num_key_value_heads * self.head_dim  # v_size
                ]

            self.q_proj = torch.nn.Linear(self.hidden_size, split_sizes[0])
            self.k_proj = torch.nn.Linear(self.hidden_size, split_sizes[1])
            self.v_proj = torch.nn.Linear(self.hidden_size, split_sizes[2])
            if config.model_type == 'chatglm':
                # chatglm-6b
                qkv_weight = self.qkv_proj.weight.data.view(self.num_heads, 3, self.head_dim, self.hidden_size)
                self.q_proj.weight.data = qkv_weight[:, 0, :, :].reshape(self.hidden_size, self.hidden_size)
                self.k_proj.weight.data = qkv_weight[:, 1, :, :].reshape(self.hidden_size, self.hidden_size)
                self.v_proj.weight.data = qkv_weight[:, 2, :, :].reshape(self.hidden_size, self.hidden_size)
                qkv_bias = self.qkv_proj.bias.data.view(self.num_heads, 3, self.head_dim)
                self.q_proj.bias.data = qkv_bias[:, 0, :].reshape(self.hidden_size)
                self.k_proj.bias.data = qkv_bias[:, 1, :].reshape(self.hidden_size)
                self.v_proj.bias.data = qkv_bias[:, 2, :].reshape(self.hidden_size)
            else:
                # other
                qw, kw, vw = torch.split(self.qkv_proj.weight, split_sizes)
                self.q_proj.weight.data = qw
                self.k_proj.weight.data = kw
                self.v_proj.weight.data = vw
                if self.qkv_proj.bias is not None:
                    qb, kb, vb = torch.split(self.qkv_proj.bias, split_sizes)
                    self.q_proj.bias.data = qb
                    self.k_proj.bias.data = kb
                    self.v_proj.bias.data = vb
                else:
                    data_type = self.q_proj.weight.dtype
                    self.q_proj.bias.data = torch.zeros(split_sizes[0], dtype=data_type)
                    self.k_proj.bias.data = torch.zeros(split_sizes[1], dtype=data_type)
                    self.v_proj.bias.data = torch.zeros(split_sizes[2], dtype=data_type)
            self.q_proj.weight.requires_grad = False
            self.k_proj.weight.requires_grad = False
            self.v_proj.weight.requires_grad = False
            self.q_proj.bias.requires_grad = False
            self.k_proj.bias.requires_grad = False
            self.v_proj.bias.requires_grad = False


    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        rotary_pos_emb: Optional[torch.Tensor] = None,
        cross_attention_states: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        bsz, q_len, _ = hidden_states.size()
        query_states = self.q_proj(hidden_states)
        if cross_attention_states is not None:
            hidden_states = cross_attention_states
        key_states = self.k_proj(hidden_states)
        value_states = self.v_proj(hidden_states)
        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim)
        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim)
        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim)
        # openelm model has qk_norm
        if hasattr(self, 'q_norm') and self.q_norm is not None and \
           hasattr(self, 'k_norm') and self.k_norm is not None :
            query_states = self.q_norm(query_states)
            key_states = self.k_norm(key_states)

        kv_seq_len = key_states.shape[1]
        if past_key_value is not None:
            kv_seq_len += past_key_value[0].shape[1]
        # rope
        if self.rotary is not None:
            cos, sin = rotary_pos_emb[0], rotary_pos_emb[1]
            query_states = self.rotary.apply_rotary_pos(query_states, cos, sin)
            key_states = self.rotary.apply_rotary_pos(key_states, cos, sin)

        if self.export_fused_attn:
            attn_output = self.fused_attn(query_states, key_states, value_states, attention_mask)
            attn_output = self.o_proj(attn_output)
            return attn_output, past_key_value

        # kv cache
        if past_key_value is not None:
            past_key, past_value = past_key_value[0], past_key_value[1]
            key_states = torch.cat((past_key, key_states), dim=1)
            value_states = torch.cat((past_value, value_states), dim=1)

        past_key_value = torch.stack((key_states, value_states))
        query_states = query_states.transpose(1, 2)
        key_states = key_states.permute([0, 2, 3, 1])
        value_states = value_states.transpose(1, 2)
        # repeat k/v heads if n_kv_heads < n_heads
        key_states = repeat_kv(key_states, self.num_key_value_groups)
        value_states = repeat_kv(value_states, self.num_key_value_groups)
        #------- attention ----------
        # query_states @ key_states
        attn_weights = torch.matmul(query_states, key_states) / math.sqrt(self.head_dim)
        # attention_mask
        if attention_mask.dtype in (torch.bool, torch.int32):
            # chatglm
            attn_weights.masked_fill_(attention_mask, -10000.0)
        else:
            attn_weights = attn_weights + attention_mask

        if hasattr(self, 'sinks'):
            sinks = self.sinks.reshape(1, -1, 1, 1).to(torch.float32).expand(query_states.shape[0], -1, query_states.shape[-2], -1)
            combined_logits = torch.cat([attn_weights, sinks], dim=-1)
            combined_logits = combined_logits - combined_logits.max(dim=-1, keepdim=True).values
            probs = torch.nn.functional.softmax(combined_logits, dim=-1, dtype=torch.float32).to(query_states.dtype)
            attn_weights = probs[..., :-1]  # we drop the sink here
        else:
            # upcast softmax to fp32
            attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
        # attn_weights @ value_states
        attn_output = torch.matmul(attn_weights, value_states)

        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.reshape(bsz, q_len, -1)
        attn_output = self.o_proj(attn_output)
        return attn_output, past_key_value

def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

def _compute_yarn_parameters(rotary_dim, base_theta, scaling_config, max_position_embeddings):
    """
    计算 YaRN (Yet another RoPE extensioN method) 的参数。
    此函数等价于 Hugging Face Transformers 中的 YaRN 实现。

    Args:
        rotary_dim (int): RoPE 的维度。
        base_theta (float): RoPE 的基础 theta 值。
        scaling_config (dict): 包含 YaRN 特定配置的字典。
        max_position_embeddings (int): 模型的最大位置编码。

    Returns:
        tuple[torch.Tensor, float]:
            - inv_freq (torch.Tensor): 计算好的、用于 RoPE 的逆频率 (即 theta)。
            - attention_scaling (float): 应用于 Query 向量的缩放因子。
    """
    def get_mscale(scale, m_scale):
        if scale <= 1:
            return 1.0
        return 0.1 * m_scale * math.log(scale) + 1.0

    def find_correction_dim(num_rotations, d, b, max_pos):
        return (d * math.log(max_pos / (num_rotations * 2 * math.pi))) / (2 * math.log(b))

    def find_correction_range(low_rot, high_rot, d, b, max_pos):
        low = find_correction_dim(low_rot, d, b, max_pos)
        high = find_correction_dim(high_rot, d, b, max_pos)
        return max(0, math.floor(low)), min(d - 1, math.ceil(high))

    def linear_ramp_factor(mn, mx, d):
        if mn == mx:
            mx += 0.001
        linear_func = (torch.arange(d, dtype=torch.float32) - mn) / (mx - mn)
        return torch.clamp(linear_func, 0, 1)

    # 1. 提取 YaRN 参数
    factor = scaling_config['factor']
    beta_fast = scaling_config.get("beta_fast", 32)
    beta_slow = scaling_config.get("beta_slow", 1)
    original_max_pos = scaling_config.get("original_max_position_embeddings", max_position_embeddings)
    mscale = scaling_config.get("mscale", 1.0)

    # 2. 计算 attention_scaling (即 attention_factor)
    attention_scaling = get_mscale(factor, mscale)

    # 3. 计算 inv_freq (即 theta)
    dim = rotary_dim

    # 计算插值和外推的频率
    pos_freqs = base_theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)
    inv_freq_extrapolation = 1.0 / pos_freqs
    inv_freq_interpolation = 1.0 / (factor * pos_freqs)

    # 找到需要修正的维度范围
    low, high = find_correction_range(beta_fast, beta_slow, dim, base_theta, original_max_pos)

    # 创建维度混合的 ramp (作用于 dim//2 的频率上)
    inv_freq_extrapolation_factor = 1 - linear_ramp_factor(low, high, dim // 2)

    # 混合插值和外推频率，得到最终的 inv_freq
    inv_freq = (
        inv_freq_interpolation * (1 - inv_freq_extrapolation_factor)
        + inv_freq_extrapolation * inv_freq_extrapolation_factor
    )

    return inv_freq, attention_scaling

class Rotary(torch.nn.Module):
    def __init__(self, config):
        super().__init__()
        if config is None: return
        self.rope_theta = config.rope_theta
        self.rope_ratio = config.rope_ratio
        self.rope_theta *= self.rope_ratio
        self.rotary_dim = config.head_dim
        self.model_type = config.model_type
        if hasattr(config, 'rotary_dim'):
            self.rotary_dim = config.rotary_dim
        if self.model_type == 'chatglm':
            self.rotary_dim = config.head_dim // 2
        self.mrope_section = None
        self.theta_sections = None
        self.attention_scaling = 1.0
        self.is_scaled = False

        def get_theta():
            return 1.0 / (self.rope_theta ** (torch.arange(0, self.rotary_dim, 2, dtype=torch.float32) / self.rotary_dim))
        # default rope type's theta
        self.theta = get_theta()
        # other type
        if hasattr(config, 'rope_scaling') and config.rope_scaling is not None:
            scaling_config = config.rope_scaling
            # get rope_type
            rope_type = 'default'
            if 'type' in config.rope_scaling:
                rope_type = config.rope_scaling['type']
            elif 'rope_type' in config.rope_scaling:
                rope_type = config.rope_scaling['rope_type']
            # gen theta for rope_type
            if rope_type == 'dynamic': # NTK
                if 'alpha' in config.rope_scaling: # NTKAlpha in Hunyuan
                    self.rope_theta *= (config.rope_scaling['alpha'] ** (self.rotary_dim / (self.rotary_dim - 2)))
                else: # NTKScaling
                    pass
                self.theta = get_theta()
            elif rope_type == 'yarn':
                self.is_scaled = True
                self.theta, self.attention_scaling = _compute_yarn_parameters(
                    rotary_dim=self.rotary_dim,
                    base_theta=self.rope_theta,
                    scaling_config=scaling_config,
                    max_position_embeddings=config.max_position_embeddings
                )
            elif rope_type == 'longrope': # longrope in MiniCPM
                self.is_scaled = True
                original_max_position_embeddings = config.rope_scaling['original_max_position_embeddings']
                scale = (config.max_position_embeddings / original_max_position_embeddings)
                self.attention_scaling = math.sqrt(1 + math.log(scale) / math.log(original_max_position_embeddings))
                # long_factor = config.rope_scaling['long_factor']
                short_factor = config.rope_scaling['short_factor']
                self.theta = get_theta() / torch.tensor(short_factor, dtype=torch.float32)

            # mrope for multimode
            if 'mrope_section' in scaling_config:
                self.mrope_section = scaling_config['mrope_section']
                self.theta_sections = get_theta().unsqueeze(0).split(self.mrope_section, dim=-1)

    def forward(self, position_ids):
        if self.theta_sections is not None:
            return self.mrope_forward(position_ids)
        position_ids = position_ids.float().reshape(-1, 1)
        idx_theta = position_ids * self.theta
        rotary_pos_emb = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)])
        if self.model_type == 'ernie4_5':
            rotary_pos_emb = torch.stack((rotary_pos_emb, rotary_pos_emb), dim=-1)
            rotary_pos_emb = rotary_pos_emb.reshape(*rotary_pos_emb.shape[:-2], -1)
        elif self.model_type != 'chatglm2':
            rotary_pos_emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
        rotary_pos_emb = rotary_pos_emb.unsqueeze(2).unsqueeze(1)
        if self.is_scaled:
            rotary_pos_emb *= self.attention_scaling
        return rotary_pos_emb

    def mrope_forward(self, position_ids):
        position_ids = position_ids.float().unsqueeze(-1)
        idx_theta = torch.concat([
            position_ids[0] * self.theta_sections[0],
            position_ids[1] * self.theta_sections[1],
            position_ids[2] * self.theta_sections[2]
        ], dim=-1)
        rotary_pos_emb = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)])
        rotary_pos_emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
        rotary_pos_emb = rotary_pos_emb.unsqueeze(2).unsqueeze(1)
        return rotary_pos_emb

    def apply_rotary_pos(self, x, cos, sin):
        if self.model_type == 'chatglm':
            return self.chatglm_rotary_pos(x, cos, sin)
        if self.model_type == 'chatglm2':
            return self.chatglm2_rotary_pos(x, cos, sin)
        if self.model_type == 'phi-msft':
            return self.phi_rotary_pos(x, cos, sin)
        if self.model_type == 'ernie4_5':
            return self.ernie_rotary_pos(x, cos, sin)
        return self.llama_rotary_pos(x, cos, sin)

    def llama_rotary_pos(self, x, cos, sin):
        x = (x * cos) + (rotate_half(x) * sin)
        return x

    def ernie_rotary_pos(self, x, cos, sin):
        rotate_half_x = torch.stack(
            [-x[:, :, :, 1::2], x[:, :, :, 0::2]], dim=-1
        ).reshape(x.shape)
        x = (x * cos) + (rotate_half_x * sin)
        return x

    def phi_rotary_pos(self, x, cos, sin):
        x, x_pass = x[..., :self.rotary_dim], x[..., self.rotary_dim:]
        x = (x * cos) + (rotate_half(x) * sin)
        return torch.cat((x, x_pass), dim=-1)

    def chatglm2_rotary_pos(self, x, cos, sin):
        x, x_pass = x[..., :self.rotary_dim], x[..., self.rotary_dim:]
        b, s, n, h = x.shape
        xshaped = x.view(b, s, n, h//2, 2)
        x = torch.concat(
            [
                xshaped[..., 0] * cos - xshaped[..., 1] * sin,
                xshaped[..., 1] * cos + xshaped[..., 0] * sin,
            ],
            -1,
        )
        return torch.cat((x, x_pass), dim=-1)

    def chatglm_rotary_pos(self, x, cos, sin):
        seq = x.shape[1]
        x1, x2 = x[..., :self.rotary_dim], x[..., self.rotary_dim:]
        cos1, sin1 = cos[:, :seq, ...], sin[:, :seq, ...]
        cos2, sin2 = cos[:, seq:, ...], sin[:, seq:, ...]
        x1 = (x1 * cos1) + (rotate_half(x1) * sin1)
        x2 = (x2 * cos2) + (rotate_half(x2) * sin2)
        return torch.cat((x1, x2), dim=-1)

class VisionRotary(Rotary):
    def __init__(self, config):
        super().__init__(config)

    # support [h_pos, w_pos]
    def forward(self, position_ids):
        # [2, patch_len, 1]
        position_ids = position_ids.float().unsqueeze(-1)
        idx_theta = position_ids * self.theta
        # [patch_len, rotary_dim]
        idx_theta = idx_theta.permute(1, 0, 2).reshape(-1, self.rotary_dim)
        rotary_pos_emb = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)])
        rotary_pos_emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
        rotary_pos_emb = rotary_pos_emb.unsqueeze(2).unsqueeze(1)
        return rotary_pos_emb

class GptOssExpert(torch.nn.Module):
    def __init__(self, hidden_size, expert_dim):
        super().__init__()
        self.expert_dim = expert_dim
        self.gate_up_proj_linear = torch.nn.Linear(hidden_size, 2 * expert_dim)
        self.down_proj_linear = torch.nn.Linear(expert_dim, hidden_size)
        self.alpha = 1.702
        self.limit = 7.0

    def forward(self, hidden_states: torch.Tensor, debug=False) -> torch.Tensor:
        gate_up = self.gate_up_proj_linear(hidden_states)
        gate, up = gate_up[..., ::2], gate_up[..., 1::2]
        # gate = gate.clamp(min=None, max=self.limit)
        limit_tensor = torch.tensor(self.limit, device=gate.device, dtype=gate.dtype)
        gate = torch.min(gate, limit_tensor)
        up = up.clamp(min=-self.limit, max=self.limit)
        glu = gate * torch.sigmoid(gate * self.alpha)
        gated_output = (up + 1) * glu
        out = self.down_proj_linear(gated_output)
        return out

class Mlp(torch.nn.Module):
    def __init__(self, mlp, config, layer_id):
        super().__init__()
        self.layer_id = layer_id
        ModelMapper.do_map(self, mlp, config.model_map['mlp'])
        self.is_moe = hasattr(self, 'experts')
        self.export_moe = False
        self.custom_moe = MoE(self.num_experts, self.top_k, layer_id)
        self.moe_type = 'qwen3_moe'
        if hasattr(self, 'router'):
            self.moe_type = 'gpt_oss'
            hidden_dim = self.router.weight.shape[1]
            self.gate = torch.nn.Linear(hidden_dim, self.num_experts, bias=True)
            self.gate.weight.data = self.router.weight.data
            self.gate.bias.data = self.router.bias.data
            # refacte experts to qwen3_experts
            original_experts = self.experts
            expert_dim = original_experts.expert_dim
            new_experts_list = torch.nn.ModuleList()
            for i in range(self.num_experts):
                expert_mlp = GptOssExpert(hidden_dim, expert_dim)
                expert_mlp.gate_up_proj_linear.weight.data = original_experts.gate_up_proj.data[i].transpose(0, 1)
                expert_mlp.gate_up_proj_linear.bias.data = original_experts.gate_up_proj_bias.data[i]
                expert_mlp.down_proj_linear.weight.data = original_experts.down_proj.data[i].transpose(0, 1)
                expert_mlp.down_proj_linear.bias.data = original_experts.down_proj_bias.data[i]
                new_experts_list.append(expert_mlp)
            self.experts = new_experts_list
            del self.router

    def forward(self, hidden_states: torch.Tensor):
        if not self.is_moe:
            # general Mlp
            return self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))

        # MoE Mlp
        batch_size, sequence_length, hidden_dim = hidden_states.shape
        hidden_states = hidden_states.view(-1, hidden_dim)
        # router_logits: (batch * sequence_length, n_experts)
        router_logits = self.gate(hidden_states)

        if self.moe_type == 'gpt_oss':
            routing_weights, selected_experts = torch.topk(router_logits, self.top_k, dim=-1)
            routing_weights = torch.nn.functional.softmax(routing_weights, dim=-1, dtype=torch.float).to(hidden_states.dtype)
        else:
            routing_weights = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float)
            routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
            if self.norm_topk_prob:  # only diff with mixtral sparse moe block!
                routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
            # we cast back to the input dtype
            routing_weights = routing_weights.to(hidden_states.dtype)

        if self.export_moe:
            return self.custom_moe(hidden_states, routing_weights, selected_experts)

        final_hidden_states = torch.zeros(
            (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
        )
        if False: # cpp impl
            seqlen, topk = selected_experts.shape
            if seqlen == 1:
                expert_idx = int(selected_experts[0, 0])
                scale = float(routing_weights[0, 0])
                output = self.experts[expert_idx](hidden_states) * scale
                for i in range(1, topk):
                    expert_idx = int(selected_experts[0, i])
                    scale = float(routing_weights[0, i])
                    output += self.experts[expert_idx](hidden_states) * scale
                return output

            hss = torch.split(hidden_states, 1)
            expertWorks = [[] for i in range(self.num_experts)]

            for i in range(seqlen):
                for j in range(topk):
                    expert_idx = int(selected_experts[i, j])
                    scale = float(routing_weights[i, j])
                    expertWorks[expert_idx].append((i, scale))

            for i in range(self.num_experts):
                if len(expertWorks[i]) == 0:
                    continue
                input_hs = []
                for token_id, scale in expertWorks[i]:
                    input_hs.append(hss[token_id])
                output_hs = self.experts[i](torch.concat(input_hs))
                output_hss = torch.split(output_hs, 1)
                for j in range(len(expertWorks[i])):
                    token_id, scale = expertWorks[i][j]
                    scale_hs = output_hss[j] * scale
                    final_hidden_states[token_id] += scale_hs.squeeze(0)
            return final_hidden_states

        # One hot encode the selected experts to create an expert mask
        # this will be used to easily index which expert is going to be sollicitated
        expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)

        # Loop over all available experts in the model and perform the computation on each expert
        for expert_idx in range(self.num_experts):
            expert_layer = self.experts[expert_idx]
            idx, top_x = torch.where(expert_mask[expert_idx])

            # Index the correct hidden states and compute the expert hidden state for
            # the current expert. We need to make sure to multiply the output hidden
            # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
            current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
            current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None]
            # However `index_add_` only support torch tensors for indexing so we'll use
            # the `top_x` tensor here.
            final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
        final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)

        return final_hidden_states

class Decoder(torch.nn.Module):
    def __init__(self, decoder, layer_id, config):
        super().__init__()
        self.cross_decoder = False
        ModelMapper.do_map(self, decoder, config.model_map['decoder'])
        if 'mlp' in config.model_map:
            self.mlp = Mlp(self.mlp, config, layer_id)
        # mllama has cross_attn
        if hasattr(self, 'cross_attn') and self.cross_attn is not None:
            self.cross_decoder = True
            self.self_attn = Attention(self.cross_attn, layer_id, config)
        else:
            self.self_attn = Attention(self.self_attn, layer_id, config)
        self.hidden_size = config.hidden_size
        if hasattr(config, 'num_hidden_layers'):
            # minicpm
            self.num_hidden_layers = config.num_hidden_layers
            # chatglm
            self.alpha = (2 * config.num_hidden_layers) ** 0.5 if config.model_type == 'chatglm' else 1.0

    def forward(
        self,
        hidden_states: torch.Tensor,
        rotary_pos_emb: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        cross_attention_states: Optional[torch.Tensor] = None,
        cross_attention_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
        hidden_states = hidden_states.view(1, -1, self.hidden_size)
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        norm_hidden_states = hidden_states
        # Self Attention
        hidden_states, present_key_value = self.self_attn(
            hidden_states=hidden_states,
            rotary_pos_emb=rotary_pos_emb,
            attention_mask=attention_mask,
            past_key_value=past_key_value,
            cross_attention_states=cross_attention_states,
        )
        # Fully Connected
        if not hasattr(self, 'post_attention_layernorm'):
            # phi
            feed_forward_hidden_states = self.mlp(norm_hidden_states)
            hidden_states = hidden_states + feed_forward_hidden_states + residual
        elif hasattr(self, 'alpha') and self.alpha != 1.0:
            # chatglm-6b
            hidden_states = norm_hidden_states * self.alpha + hidden_states
            mlp_input = self.post_attention_layernorm(hidden_states)
            mlp_output = self.mlp(mlp_input)
            hidden_states = mlp_input * self.alpha + mlp_output
        elif hasattr(self, 'pre_feedforward_layernorm'):
            # gemma2
            hidden_states = self.post_attention_layernorm(hidden_states)
            hidden_states = residual + hidden_states
            residual = hidden_states
            hidden_states = self.pre_feedforward_layernorm(hidden_states)
            hidden_states = self.mlp(hidden_states)
            hidden_states = self.post_feedforward_layernorm(hidden_states)
            hidden_states = residual + hidden_states
        elif cross_attention_mask is not None and self.cross_decoder:
            hidden_states = residual + self.cross_attn_attn_gate.tanh() * hidden_states
            residual = hidden_states
            hidden_states = self.post_attention_layernorm(hidden_states)
            hidden_states = self.mlp(hidden_states)
            hidden_states = cross_attention_mask * hidden_states
            hidden_states = residual + self.cross_attn_mlp_gate.tanh() * hidden_states
        elif hasattr(self, 'scale_depth'):
            # minicpm
            hidden_states = residual + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers))
            residual = hidden_states
            hidden_states = self.post_attention_layernorm(hidden_states)
            hidden_states = self.mlp(hidden_states)
            hidden_states = residual + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers))
        else:
            # general
            hidden_states = residual + hidden_states
            residual = hidden_states
            hidden_states = self.post_attention_layernorm(hidden_states)
            hidden_states = self.mlp(hidden_states)
            hidden_states = residual + hidden_states

        return hidden_states, present_key_value

class Lm(torch.nn.Module):
    def __init__(self, lm_):
        super().__init__()
        self.lm = lm_

    def forward(self, hidden_states):
        m_logits = self.lm(hidden_states)
        return m_logits