"""Eather-500M: custom decoder-only Transformer in pure PyTorch.
RoPE, grouped-query attention, RMSNorm, SwiGLU, tied token embeddings.
"""
import json
from dataclasses import dataclass
from pathlib import Path
import torch
from torch import nn
import torch.nn.functional as F


@dataclass
class ModelConfig:
    vocab_size: int = 32000
    d_model: int = 1280
    n_layers: int = 26
    n_heads: int = 20
    n_kv_heads: int = 5
    ffn_hidden: int = 3584
    max_seq_len: int = 2048
    rope_theta: float = 10000.0
    dropout: float = 0.0
    tie_embeddings: bool = True

    @classmethod
    def load(cls, path):
        return cls(**json.loads(Path(path).read_text()))


def rms_norm(x, weight, eps=1e-6):
    # Accumulate normalization in fp32 for stability.
    var = x.float().pow(2).mean(dim=-1, keepdim=True)
    return (x.float() * torch.rsqrt(var + eps)).to(x.dtype) * weight


class RMSNorm(nn.Module):
    def __init__(self, size):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(size))

    def forward(self, x):
        return rms_norm(x, self.weight)


def rope(q, k, theta, offset=0):
    d = q.shape[-1]
    positions = torch.arange(offset, offset + q.shape[-2], device=q.device, dtype=torch.float32)
    freqs = 1.0 / (theta ** (torch.arange(0, d, 2, device=q.device, dtype=torch.float32) / d))
    angles = torch.outer(positions, freqs)
    cos = angles.cos()[None, None, :, :]
    sin = angles.sin()[None, None, :, :]

    def rotate(x):
        a, b = x.float()[..., ::2], x.float()[..., 1::2]
        return torch.stack((a * cos - b * sin, a * sin + b * cos), dim=-1).flatten(-2).to(x.dtype)

    return rotate(q), rotate(k)


class Attention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        assert cfg.d_model % cfg.n_heads == 0
        assert cfg.n_heads % cfg.n_kv_heads == 0
        self.head_dim = cfg.d_model // cfg.n_heads
        self.n_heads = cfg.n_heads
        self.n_kv_heads = cfg.n_kv_heads
        self.theta = cfg.rope_theta
        self.dropout = cfg.dropout
        self.q = nn.Linear(cfg.d_model, cfg.n_heads * self.head_dim, bias=False)
        self.k = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False)
        self.v = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False)
        self.o = nn.Linear(cfg.d_model, cfg.d_model, bias=False)

    def forward(self, x):
        b, t, _ = x.shape
        q = self.q(x).view(b, t, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k(x).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v(x).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2)
        q, k = rope(q, k, self.theta)
        # Explicit head repeats for compatibility with PyTorch versions without enable_gqa.
        repeats = self.n_heads // self.n_kv_heads
        k = k.repeat_interleave(repeats, dim=1)
        v = v.repeat_interleave(repeats, dim=1)
        out = F.scaled_dot_product_attention(
            q, k, v, is_causal=True,
            dropout_p=self.dropout if self.training else 0.0,
        )
        out = out.transpose(1, 2).contiguous().view(b, t, -1)
        return self.o(out)


class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.d_model)
        self.attn = Attention(cfg)
        self.ff_norm = RMSNorm(cfg.d_model)
        self.up = nn.Linear(cfg.d_model, cfg.ffn_hidden, bias=False)
        self.gate = nn.Linear(cfg.d_model, cfg.ffn_hidden, bias=False)
        self.down = nn.Linear(cfg.ffn_hidden, cfg.d_model, bias=False)

    def forward(self, x):
        x = x + self.attn(self.attn_norm(x))
        y = self.ff_norm(x)
        return x + self.down(F.silu(self.gate(y)) * self.up(y))


class Eather500M(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.cfg = cfg
        self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layers))
        self.norm = RMSNorm(cfg.d_model)
        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        if cfg.tie_embeddings:
            self.lm_head.weight = self.embed.weight
        self.apply(self._init_weights)
        # Small residual projections improve deep-transformer stability.
        for layer in self.blocks:
            nn.init.normal_(layer.attn.o.weight, mean=0, std=0.02 / (2 * cfg.n_layers) ** 0.5)
            nn.init.normal_(layer.down.weight, mean=0, std=0.02 / (2 * cfg.n_layers) ** 0.5)
        self.gradient_checkpointing = False

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

    def forward(self, input_ids, labels=None, loss_mask=None):
        if input_ids.shape[1] > self.cfg.max_seq_len:
            raise ValueError(f"Context exceeds {self.cfg.max_seq_len} tokens")
        x = self.embed(input_ids)
        for block in self.blocks:
            if self.gradient_checkpointing and self.training:
                x = torch.utils.checkpoint.checkpoint(block, x, use_reentrant=False)
            else:
                x = block(x)
        logits = self.lm_head(self.norm(x))
        loss = None
        if labels is not None:
            # Shift inputs and labels in the dataset, so no shift happens here.
            per_token = F.cross_entropy(
                logits.float().reshape(-1, logits.size(-1)), labels.reshape(-1),
                reduction="none",
            )
            if loss_mask is not None:
                mask = loss_mask.reshape(-1).float()
                loss = (per_token * mask).sum() / mask.sum().clamp_min(1)
            else:
                loss = per_token.mean()
        return logits, loss


if __name__ == "__main__":
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument("--config", default="config.json")
    p.add_argument("--device", default="meta", help="Use 'meta' to count without allocating RAM")
    args = p.parse_args()
    with torch.device(args.device):
        model = Eather500M(ModelConfig.load(args.config))
    n = sum(p.numel() for p in model.parameters())
    print(f"Parameters: {n:,} ({n / 1e6:.2f}M)")