import math
from dataclasses import dataclass
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_laoyao_moe import LaoyaoMoeConfig


class LaoyaoRMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y = x.float() * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return (y * self.weight.float()).to(dtype=x.dtype)


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = torch.chunk(x, 2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def apply_rope(q: torch.Tensor, k: torch.Tensor, rope_theta: float) -> tuple[torch.Tensor, torch.Tensor]:
    seq_len = q.shape[1]
    head_dim = q.shape[-1]
    inv_freq = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2, device=q.device, dtype=torch.float32) / head_dim))
    pos = torch.arange(seq_len, device=q.device, dtype=torch.float32)
    freqs = torch.einsum("t,d->td", pos, inv_freq)
    emb = torch.cat((freqs, freqs), dim=-1)
    cos = emb.cos()[None, :, None, :]
    sin = emb.sin()[None, :, None, :]
    qf = q.float()
    kf = k.float()
    return (qf * cos + rotate_half(qf) * sin).to(q.dtype), (kf * cos + rotate_half(kf) * sin).to(k.dtype)


class LaoyaoAttention(nn.Module):
    def __init__(self, config: LaoyaoMoeConfig):
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.hidden_size // config.num_attention_heads
        self.hidden_size = config.hidden_size
        self.rope_theta = config.rope_theta
        qkv_out = (self.num_heads + 2 * self.num_kv_heads) * self.head_dim
        self.qkv_proj = nn.Linear(config.hidden_size, qkv_out, bias=False)
        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        qkv = self.qkv_proj(x)
        bsz, seq_len, _ = x.shape
        heads_per_group = self.num_heads // self.num_kv_heads
        qkv = qkv.view(bsz, seq_len, self.num_kv_heads, heads_per_group + 2, self.head_dim)
        q = qkv[:, :, :, :heads_per_group, :].reshape(bsz, seq_len, self.num_heads, self.head_dim)
        k = qkv[:, :, :, heads_per_group, :]
        v = qkv[:, :, :, heads_per_group + 1, :]
        q, k = apply_rope(q, k, self.rope_theta)
        repeat = self.num_heads // self.num_kv_heads
        k = k.repeat_interleave(repeat, dim=2)
        v = v.repeat_interleave(repeat, dim=2)
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        attn = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.head_dim)
        mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool), diagonal=1)
        attn = attn.masked_fill(mask[None, None, :, :], torch.finfo(attn.dtype).min)
        out = torch.matmul(torch.softmax(attn, dim=-1), v)
        out = out.transpose(1, 2).contiguous().view(bsz, seq_len, self.hidden_size).to(x.dtype)
        return self.o_proj(out)


class LaoyaoDenseMLP(nn.Module):
    def __init__(self, config: LaoyaoMoeConfig):
        super().__init__()
        self.fc1 = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)
        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        a, b = self.fc1(x).chunk(2, dim=-1)
        return self.fc2(F.silu(a) * b)


class LaoyaoMoeMLP(nn.Module):
    def __init__(self, config: LaoyaoMoeConfig):
        super().__init__()
        self.num_experts = config.num_experts
        self.top_k = config.num_experts_per_tok
        self.capacity_factor = config.moe_expert_capacity_factor
        self.router = nn.Linear(config.hidden_size, config.num_experts, bias=False)
        self.experts_fc1 = nn.Parameter(torch.empty(config.num_experts, 2 * config.moe_intermediate_size, config.hidden_size))
        self.experts_fc2 = nn.Parameter(torch.empty(config.num_experts, config.hidden_size, config.moe_intermediate_size))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        flat_x = x.reshape(-1, x.shape[-1])
        router_logits = F.linear(flat_x, self.router.weight)
        topv, topi = torch.topk(router_logits, self.top_k, dim=-1, sorted=torch.is_grad_enabled())
        probs = torch.softmax(topv, dim=-1)
        routing_probs = torch.zeros_like(router_logits).scatter(1, topi, probs.type_as(router_logits))
        routing_map = torch.zeros_like(router_logits, dtype=torch.int32).scatter(1, topi, 1).bool()
        if self.capacity_factor is not None:
            num_tokens = routing_probs.shape[0]
            capacity = math.ceil((num_tokens * self.top_k / self.num_experts) * self.capacity_factor)
            if capacity <= num_tokens:
                _, keep = torch.topk(routing_probs, k=capacity, dim=0, sorted=False)
                capacity_mask = torch.zeros_like(routing_probs, dtype=torch.int32).scatter(0, keep, 1).bool()
                routing_map = routing_map & capacity_mask
                routing_probs = routing_probs * routing_map
        flat_out = torch.zeros_like(flat_x)
        for expert_id in range(self.num_experts):
            mask = routing_map[:, expert_id]
            if not mask.any():
                continue
            token_idx = mask.nonzero(as_tuple=True)[0]
            expert_in = flat_x[token_idx]
            hidden = F.linear(expert_in, self.experts_fc1[expert_id])
            a, b = hidden.chunk(2, dim=-1)
            hidden = F.silu(a) * b
            expert_out = F.linear(hidden, self.experts_fc2[expert_id])
            flat_out[token_idx] += expert_out * routing_probs[token_idx, expert_id].unsqueeze(-1)
        return flat_out.reshape_as(x).to(x.dtype)


class LaoyaoDecoderLayer(nn.Module):
    def __init__(self, config: LaoyaoMoeConfig, layer_idx: int):
        super().__init__()
        self.input_layernorm = LaoyaoRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.self_attn = LaoyaoAttention(config)
        self.is_moe = layer_idx in set(config.moe_layers)
        self.pre_mlp_layernorm = LaoyaoRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.mlp = LaoyaoMoeMLP(config) if self.is_moe else LaoyaoDenseMLP(config)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.self_attn(self.input_layernorm(x))
        x = x + self.mlp(self.pre_mlp_layernorm(x))
        return x


class LaoyaoMoePreTrainedModel(PreTrainedModel):
    config_class = LaoyaoMoeConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False


class LaoyaoMoeModel(LaoyaoMoePreTrainedModel):
    def __init__(self, config: LaoyaoMoeConfig):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([LaoyaoDecoderLayer(config, idx) for idx in range(config.num_hidden_layers)])
        self.norm = LaoyaoRMSNorm(config.hidden_size, config.rms_norm_eps)

    def forward(self, input_ids: torch.LongTensor) -> torch.Tensor:
        x = self.embed_tokens(input_ids)
        for layer in self.layers:
            x = layer(x)
        return self.norm(x)


class LaoyaoMoeForCausalLM(LaoyaoMoePreTrainedModel, GenerationMixin):
    _tied_weights_keys = ["lm_head.weight"]
    all_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}

    def __init__(self, config: LaoyaoMoeConfig):
        super().__init__(config)
        self.config.use_cache = False
        self.model = LaoyaoMoeModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.tie_weights()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, value):
        self.lm_head = value

    def forward(
        self,
        input_ids: torch.LongTensor,
        labels: Optional[torch.LongTensor] = None,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        hidden = self.model(input_ids)
        logits = self.lm_head(hidden).float()
        loss = None
        if labels is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
        return CausalLMOutputWithPast(loss=loss, logits=logits)
