import math

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

try:
    from .configuration_mini_llm import MiniLlmConfig
except ImportError:
    from configuration_mini_llm import MiniLlmConfig


class MiniLlmRMSNorm(nn.Module):
    def __init__(self, d, eps=1e-5):
        super().__init__()
        self.w = nn.Parameter(torch.ones(d))
        self.eps = eps

    def forward(self, x):
        n = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
        return (x.float() * n).type_as(x) * self.w


def mini_rope_cache(T, hd, device, theta=10000.0):
    inv = 1.0 / (theta ** (torch.arange(0, hd, 2, device=device).float() / hd))
    f = torch.outer(torch.arange(T, device=device).float(), inv)
    return f.cos(), f.sin()


def mini_apply_rope(x, cos, sin):
    x1, x2 = x[..., ::2], x[..., 1::2]
    cos, sin = cos[: x.size(2)], sin[: x.size(2)]
    out = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
    return out.flatten(-2).type_as(x)


class MiniLlmAttention(nn.Module):
    def __init__(self, config: MiniLlmConfig):
        super().__init__()
        d = config.hidden_size
        self.h = config.num_attention_heads
        self.qkv = nn.Linear(d, 3 * d, bias=False)
        self.o = nn.Linear(d, d, bias=False)

    def forward(self, x, cos, sin):
        B, T, C = x.shape
        q, k, v = self.qkv(x).view(B, T, 3, self.h, C // self.h).permute(2, 0, 3, 1, 4)
        q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
        q, k = mini_apply_rope(q, cos, sin), mini_apply_rope(k, cos, sin)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.o(y.transpose(1, 2).reshape(B, T, C))


class MiniLlmSwiGLU(nn.Module):
    def __init__(self, config: MiniLlmConfig):
        super().__init__()
        d, ffn = config.hidden_size, config.intermediate_size
        self.w1 = nn.Linear(d, ffn, bias=False)
        self.w3 = nn.Linear(d, ffn, bias=False)
        self.w2 = nn.Linear(ffn, d, bias=False)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class MiniLlmBlock(nn.Module):
    def __init__(self, config: MiniLlmConfig):
        super().__init__()
        self.n1, self.n2 = MiniLlmRMSNorm(config.hidden_size, config.rms_norm_eps), MiniLlmRMSNorm(
            config.hidden_size, config.rms_norm_eps
        )
        self.attn, self.mlp = MiniLlmAttention(config), MiniLlmSwiGLU(config)

    def forward(self, x, cos, sin):
        x = x + self.attn(self.n1(x), cos, sin)
        return x + self.mlp(self.n2(x))


class MiniLlmForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = MiniLlmConfig
    base_model_prefix = "mini_llm"
    supports_gradient_checkpointing = False
    _tied_weights_keys = {"head.weight": "emb.weight"}

    def __init__(self, config: MiniLlmConfig):
        super().__init__(config)
        self.config = config
        self.emb = nn.Embedding(config.vocab_size, config.hidden_size)
        self.blocks = nn.ModuleList(MiniLlmBlock(config) for _ in range(config.num_hidden_layers))
        self.norm = MiniLlmRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.head.weight = self.emb.weight
        cos, sin = mini_rope_cache(
            config.max_position_embeddings,
            config.hidden_size // config.num_attention_heads,
            "cpu",
            config.rope_theta,
        )
        self.register_buffer("cos", cos, persistent=True)
        self.register_buffer("sin", sin, persistent=True)
        self.post_init()
        for n, p in self.named_parameters():
            if n.endswith("attn.o.weight") or n.endswith("mlp.w2.weight"):
                nn.init.normal_(p, 0.0, 0.02 / math.sqrt(2 * config.num_hidden_layers))

    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, 0.0, 0.02)

    def get_input_embeddings(self):
        return self.emb

    def set_input_embeddings(self, value):
        self.emb = value

    def get_output_embeddings(self):
        return self.head

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

    def forward(self, input_ids=None, labels=None, attention_mask=None, **kwargs):
        kwargs.pop("past_key_values", None)
        kwargs.pop("use_cache", None)
        kwargs.pop("position_ids", None)
        kwargs.pop("inputs_embeds", None)
        if input_ids is None:
            raise ValueError("input_ids is required")
        x = self.emb(input_ids)
        for b in self.blocks:
            x = b(x, self.cos, self.sin)
        logits = self.head(self.norm(x))
        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits.float().view(-1, logits.size(-1)), labels.view(-1))
        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=None)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids}


try:
    from transformers import AutoConfig, AutoModelForCausalLM

    AutoConfig.register("mini_llm", MiniLlmConfig)
    AutoModelForCausalLM.register(MiniLlmConfig, MiniLlmForCausalLM)
except Exception:
    pass