"""Decoder-only Transformer used by the educational FineWeb model family."""

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

from .configuration_fineweb import FineWebConfig


class RMSNorm(nn.Module):
    def __init__(self, width, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(width))
        self.eps = eps

    def forward(self, x):
        scale = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return x * scale * self.weight


def apply_rope(x, cos, sin):
    even = x[..., 0::2]
    odd = x[..., 1::2]
    rotated = torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1)
    return rotated.flatten(start_dim=-2)


class CausalSelfAttention(nn.Module):
    def __init__(self, width, heads, context_length):
        super().__init__()
        if width % heads:
            raise ValueError("d_model must be divisible by n_heads")
        self.heads = heads
        self.head_dim = width // heads
        self.context_length = context_length
        self.qkv = nn.Linear(width, 3 * width, bias=False)
        self.output = nn.Linear(width, width, bias=False)

    def forward(self, x):
        batch, tokens, width = x.shape
        if tokens > self.context_length:
            raise ValueError("Input exceeds configured context length")
        qkv = self.qkv(x).view(batch, tokens, 3, self.heads, self.head_dim)
        query, key, value = qkv.unbind(dim=2)
        query = query.transpose(1, 2)
        key = key.transpose(1, 2)
        value = value.transpose(1, 2)
        inv_frequency = 1.0 / (
            10000
            ** (
                torch.arange(0, self.head_dim, 2, device=x.device).float()
                / self.head_dim
            )
        )
        frequencies = torch.outer(
            torch.arange(tokens, device=x.device, dtype=torch.float), inv_frequency
        )
        cos = frequencies.cos().to(dtype=query.dtype)[None, None, :, :]
        sin = frequencies.sin().to(dtype=query.dtype)[None, None, :, :]
        query = apply_rope(query, cos, sin)
        key = apply_rope(key, cos, sin)
        attended = F.scaled_dot_product_attention(query, key, value, is_causal=True)
        attended = attended.transpose(1, 2).contiguous().view(batch, tokens, width)
        return self.output(attended)


class SwiGLU(nn.Module):
    def __init__(self, width, hidden):
        super().__init__()
        self.gate_and_up = nn.Linear(width, 2 * hidden, bias=False)
        self.down = nn.Linear(hidden, width, bias=False)

    def forward(self, x):
        gate, up = self.gate_and_up(x).chunk(2, dim=-1)
        return self.down(F.silu(gate) * up)


class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.attention_norm = RMSNorm(config.d_model, config.rms_norm_eps)
        self.attention = CausalSelfAttention(
            config.d_model, config.n_heads, config.context_length
        )
        self.mlp_norm = RMSNorm(config.d_model, config.rms_norm_eps)
        self.mlp = SwiGLU(config.d_model, config.mlp_hidden)

    def forward(self, x):
        x = x + self.attention(self.attention_norm(x))
        return x + self.mlp(self.mlp_norm(x))


class FineWebForCausalLM(PreTrainedModel, GenerationMixin):
    model_type = "fineweb_decoder"
    config_class = FineWebConfig
    base_model_prefix = "fineweb"
    main_input_name = "input_ids"
    _tied_weights_keys = (
        {"lm_head.weight": "embedding.weight"}
        if int(transformers.__version__.split(".", 1)[0]) >= 5
        else ["lm_head.weight"]
    )

    def __init__(self, config):
        super().__init__(config)
        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.blocks = nn.ModuleList(
            [TransformerBlock(config) for _ in range(config.n_layers)]
        )
        self.final_norm = RMSNorm(config.d_model, config.rms_norm_eps)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        self.post_init()

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

    def get_input_embeddings(self):
        return self.embedding

    def set_input_embeddings(self, value):
        self.embedding = 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=None, labels=None, return_dict=True, **kwargs):
        if input_ids is None:
            raise ValueError("input_ids is required")
        input_ids = input_ids[:, -self.config.context_length :]
        x = self.embedding(input_ids)
        for block in self.blocks:
            x = block(x)
        logits = self.lm_head(self.final_norm(x))
        loss = None
        if labels is not None:
            labels = labels[:, -input_ids.size(1) :]
            loss = F.cross_entropy(
                logits[:, :-1].contiguous().view(-1, self.config.vocab_size),
                labels[:, 1:].contiguous().view(-1),
                ignore_index=-100,
            )
        if not return_dict:
            return (loss, logits) if loss is not None else (logits,)
        return CausalLMOutput(loss=loss, logits=logits)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids[:, -self.config.context_length :], "use_cache": False}