"""Reference implementation of the dense decoder-only transformer behind the Standalone and Sampled standalone
models of "ShellElastic: Scalable Elastic Pretraining with Simultaneous Maximal Updates and Hyperparameter
Transfer".

Each released model is an ordinary dense decoder. The released weights already include the constant forward
multipliers of the model (for these models, the FFN output multiplier 1/e, where e is the FFN expansion), so this
file contains no scaling constants.

Architecture (one layer, pre-norm):

    h   = x + Attention(RMSNorm(x))
    out = h + SwiGLU(RMSNorm(h))

* RMSNorm with eps 1e-5; this file computes the statistics in float32 and casts the result back to the input dtype
  before the learned gain is applied.
* Attention: separate query/key/value projections without biases; queries and keys are RMS-normalized within
  each head (the gains are stored at full width, one per channel); rotary position embeddings in the
  rotate-half form with base 10,000; causal softmax attention with scale 1/sqrt(head_dim); the attention output
  is multiplied elementwise by sigmoid(x W_gate^T) before the output projection.
* SwiGLU feed-forward network: down_proj(silu(gate_proj(x)) * up_proj(x)).
* Untied input embedding and output head, no biases anywhere, context length 2,048, and the 50,432-token GPT-NeoX
  vocabulary (Table 8 of the paper; token ids of the GPT-NeoX-20B tokenizer, Section 4.1).

Rotary tables: ``rope_precision`` is ``"bfloat16"`` (the default) or ``"float32"``. ``"bfloat16"`` rounds the
positions and inverse frequencies to bfloat16, their product to bfloat16, and cos/sin of those angles to bfloat16,
in every compute precision; ``"float32"`` computes the tables in float32.

Dependencies: torch, and safetensors for loading the weights.

Example:

    import torch
    from modeling_dense import DenseDecoder

    model = DenseDecoder.from_pretrained("path/to/model_dir", device="cuda", dtype=torch.bfloat16)
    input_ids = torch.tensor([[...]], device="cuda")    # GPT-NeoX-20B token ids
    logits = model(input_ids)                             # [batch, tokens, 50432], float32
    logits, loss = model(input_ids[:, :-1], labels=input_ids[:, 1:])
"""

from __future__ import annotations

import json
import os
from dataclasses import asdict, dataclass, fields

import torch
import torch.nn.functional as F
from torch import nn

__all__ = ["DenseConfig", "DenseDecoder"]


@dataclass
class DenseConfig:
    vocab_size: int = 50432
    hidden_size: int = 384
    intermediate_size: int = 1536
    num_hidden_layers: int = 32
    num_attention_heads: int = 6
    head_dim: int = 64
    max_position_embeddings: int = 2048
    rope_theta: float = 10000.0
    rope_precision: str = "bfloat16"
    rms_norm_eps: float = 1e-5

    @classmethod
    def from_dict(cls, values: dict) -> "DenseConfig":
        """Build a config from a dict, ignoring keys that only describe the release (method, tier, ...)."""
        known = {f.name for f in fields(cls)}
        return cls(**{k: v for k, v in values.items() if k in known})

    @classmethod
    def from_json_file(cls, path: str) -> "DenseConfig":
        with open(path) as f:
            return cls.from_dict(json.load(f))

    def to_dict(self) -> dict:
        return asdict(self)


class RMSNorm(nn.Module):
    """RMSNorm over the last dimension, or within consecutive groups of ``group_size`` channels.

    The normalization runs in float32; the result is cast back to the input dtype and then multiplied by the
    per-channel gain (one gain per channel, also in grouped mode).
    """

    def __init__(self, dim: int, eps: float, group_size: int | None = None) -> None:
        super().__init__()
        self.eps = eps
        self.group_size = group_size
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = x.float()
        if self.group_size is not None:
            h = h.unflatten(-1, (-1, self.group_size))
        h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + self.eps)
        if self.group_size is not None:
            h = h.flatten(-2)
        return h.to(x.dtype) * self.weight


def rotary_cos_sin(config: DenseConfig, num_positions: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
    """Rotary tables of shape [num_positions, head_dim] in float32 (rotate-half layout)."""
    dim = config.head_dim
    frequencies = 1.0 / (config.rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))  # on CPU
    frequencies = frequencies.to(device)
    positions = torch.arange(num_positions, device=device, dtype=torch.float32)
    if config.rope_precision == "bfloat16":
        # Every step rounded to bfloat16 (the product of two bfloat16 numbers is exact in float32, so rounding it
        # once gives the bfloat16 product).
        def bf16(t: torch.Tensor) -> torch.Tensor:
            return t.to(torch.bfloat16).float()

        angles = bf16(torch.outer(bf16(positions), bf16(frequencies)))
        angles = torch.cat((angles, angles), dim=-1)
        return bf16(angles.cos()), bf16(angles.sin())
    if config.rope_precision != "float32":
        raise ValueError(f"rope_precision must be 'bfloat16' or 'float32', got {config.rope_precision!r}")
    angles = torch.outer(positions, frequencies)
    angles = torch.cat((angles, angles), dim=-1)
    return angles.cos(), angles.sin()


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """(a, b) -> (-b, a) over the two halves of the last dimension."""
    half = x.shape[-1] // 2
    return torch.cat((-x[..., half:], x[..., :half]), dim=-1)


def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """x: [batch, tokens, heads, head_dim]; cos, sin: [tokens, head_dim], already in x's dtype."""
    cos = cos[None, :, None, :]
    sin = sin[None, :, None, :]
    return x * cos + rotate_half(x) * sin


class Attention(nn.Module):
    def __init__(self, config: DenseConfig) -> None:
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.head_dim = config.head_dim
        width = config.num_attention_heads * config.head_dim
        self.q_proj = nn.Linear(config.hidden_size, width, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, width, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, width, bias=False)
        self.gate_proj = nn.Linear(config.hidden_size, width, bias=False)  # sigmoid output gate
        self.o_proj = nn.Linear(width, config.hidden_size, bias=False)
        self.q_norm = RMSNorm(width, config.rms_norm_eps, group_size=config.head_dim)
        self.k_norm = RMSNorm(width, config.rms_norm_eps, group_size=config.head_dim)

    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        batch, length, _ = x.shape
        heads = (batch, length, self.num_heads, self.head_dim)
        q = self.q_norm(self.q_proj(x)).view(heads)
        k = self.k_norm(self.k_proj(x)).view(heads)
        v = self.v_proj(x).view(heads)
        q = apply_rotary(q, cos, sin)
        k = apply_rotary(k, cos, sin)
        out = F.scaled_dot_product_attention(
            q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=True, scale=self.head_dim**-0.5
        )
        out = out.transpose(1, 2).reshape(batch, length, self.num_heads * self.head_dim)
        out = out * torch.sigmoid(self.gate_proj(x))
        return self.o_proj(out)


class MLP(nn.Module):
    def __init__(self, config: DenseConfig) -> None:
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class DecoderLayer(nn.Module):
    def __init__(self, config: DenseConfig) -> None:
        super().__init__()
        self.attn_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.attn = Attention(config)
        self.mlp_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.mlp = MLP(config)

    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        h = x + self.attn(self.attn_norm(x), cos, sin)
        return h + self.mlp(self.mlp_norm(h))


class DenseDecoder(nn.Module):
    """Decoder-only language model. ``forward`` returns float32 logits (and the mean cross-entropy loss when
    ``labels`` is given; label -100 is ignored)."""

    def __init__(self, config: DenseConfig) -> None:
        super().__init__()
        self.config = config
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList(DecoderLayer(config) for _ in range(config.num_hidden_layers))
        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

    def hidden_states(self, input_ids: torch.Tensor) -> torch.Tensor:
        """Final hidden states after the last norm, [batch, tokens, hidden_size], in the model's dtype."""
        length = input_ids.shape[1]
        if length > self.config.max_position_embeddings:
            raise ValueError(f"{length} tokens exceed the context of {self.config.max_position_embeddings}")
        x = self.embed_tokens(input_ids)
        cos, sin = rotary_cos_sin(self.config, length, input_ids.device)
        cos, sin = cos.to(x.dtype), sin.to(x.dtype)
        for layer in self.layers:
            x = layer(x, cos, sin)
        return self.norm(x)

    def forward(self, input_ids: torch.Tensor, labels: torch.Tensor | None = None):
        logits = self.lm_head(self.hidden_states(input_ids)).float()
        if labels is None:
            return logits
        loss = F.cross_entropy(logits.flatten(0, 1), labels.flatten(), ignore_index=-100)
        return logits, loss

    @torch.no_grad()
    def generate(self, input_ids: torch.Tensor, max_new_tokens: int, temperature: float = 0.0) -> torch.Tensor:
        """Simple sampling without a key/value cache (greedy when temperature is 0)."""
        for _ in range(max_new_tokens):
            context = input_ids[:, -self.config.max_position_embeddings:]
            next_logits = self(context)[:, -1]
            if temperature > 0:
                probs = torch.softmax(next_logits / temperature, dim=-1)
                next_ids = torch.multinomial(probs, num_samples=1)
            else:
                next_ids = next_logits.argmax(dim=-1, keepdim=True)
            input_ids = torch.cat([input_ids, next_ids], dim=1)
        return input_ids

    @classmethod
    def from_pretrained(cls, directory: str, device="cpu", dtype: torch.dtype = torch.float32) -> "DenseDecoder":
        """Load ``config.json`` and ``model.safetensors`` from a released model directory.

        The weights are read and cast to ``dtype`` on the CPU and then moved to ``device``, so the device holds only
        one copy of the model, in ``dtype``."""
        from safetensors.torch import load_file

        config = DenseConfig.from_json_file(os.path.join(directory, "config.json"))
        state = load_file(os.path.join(directory, "model.safetensors"), device="cpu")
        if dtype is not None:
            for name in list(state):  # cast on the CPU, one tensor at a time
                if state[name].is_floating_point():
                    state[name] = state[name].to(dtype)
        with torch.device("meta"):
            model = cls(config)
        model.load_state_dict(state, strict=True, assign=True)
        return model.to(device).eval()