import math
from dataclasses import dataclass
import torch
import torch.nn as nn
import torch.nn.functional as F

@dataclass
class ModelConfig:
    vocab_size: int = 32768
    max_seq_len: int = 512
    hidden_dim: int = 768
    intermediate_dim: int = 2048  # SwiGLU dimension
    num_layers: int = 12
    num_heads: int = 12
    eps: float = 1e-5

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

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

def precompute_rope_cis(dim: int, max_seq_len: int, theta: float = 10000.0) -> torch.Tensor:
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(max_seq_len, dtype=torch.float32)
    freqs = torch.outer(t, freqs)
    cis = torch.polar(torch.ones_like(freqs), freqs)
    return cis  # [max_seq_len, dim // 2]

def apply_rope(x: torch.Tensor, cis: torch.Tensor) -> torch.Tensor:
    B, T, H, D = x.shape
    x_complex = torch.view_as_complex(x.float().reshape(B, T, H, -1, 2))
    cis = cis[:T].unsqueeze(0).unsqueeze(2)  # Broadcast to [1, T, 1, D // 2]
    x_rotated = torch.view_as_real(x_complex * cis).reshape(B, T, H, D)
    return x_rotated.type_as(x)

class CausalSelfAttention(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.num_heads = cfg.num_heads
        self.head_dim = cfg.hidden_dim // cfg.num_heads
        self.q_proj = nn.Linear(cfg.hidden_dim, cfg.hidden_dim, bias=False)
        self.k_proj = nn.Linear(cfg.hidden_dim, cfg.hidden_dim, bias=False)
        self.v_proj = nn.Linear(cfg.hidden_dim, cfg.hidden_dim, bias=False)
        self.out_proj = nn.Linear(cfg.hidden_dim, cfg.hidden_dim, bias=False)

    def forward(self, x: torch.Tensor, cis: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape
        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim)
        k = self.k_proj(x).view(B, T, self.num_heads, self.head_dim)
        v = self.v_proj(x).view(B, T, self.num_heads, self.head_dim)

        q = apply_rope(q, cis)
        k = apply_rope(k, cis)

        q = q.transpose(1, 2)  # [B, H, T, D]
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # Dispatches automatically to FlashAttention or Cutlass
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.out_proj(out)

class SwiGLUMLP(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.w1 = nn.Linear(cfg.hidden_dim, cfg.intermediate_dim, bias=False)
        self.w2 = nn.Linear(cfg.hidden_dim, cfg.intermediate_dim, bias=False)
        self.w3 = nn.Linear(cfg.intermediate_dim, cfg.hidden_dim, bias=False)

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

class TransformerBlock(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.hidden_dim, eps=cfg.eps)
        self.attn = CausalSelfAttention(cfg)
        self.mlp_norm = RMSNorm(cfg.hidden_dim, eps=cfg.eps)
        self.mlp = SwiGLUMLP(cfg)

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

class DecoderLM(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.cfg = cfg
        self.tok_embeddings = nn.Embedding(cfg.vocab_size, cfg.hidden_dim)
        self.layers = nn.ModuleList([TransformerBlock(cfg) for _ in range(cfg.num_layers)])
        self.norm = RMSNorm(cfg.hidden_dim, eps=cfg.eps)
        self.lm_head = nn.Linear(cfg.hidden_dim, cfg.vocab_size, bias=False)

        # Weight tying
        self.lm_head.weight = self.tok_embeddings.weight

        # Precompute RoPE table
        rope_cis = precompute_rope_cis(cfg.hidden_dim // cfg.num_heads, cfg.max_seq_len)
        self.register_buffer("rope_cis", rope_cis, persistent=False)

        self.apply(self._init_weights)

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

    def forward(self, input_ids: torch.Tensor, targets: torch.Tensor | None = None):
        B, T = input_ids.shape
        x = self.tok_embeddings(input_ids)
        cis = self.rope_cis[:T]

        for layer in self.layers:
            x = layer(x, cis)
        x = self.norm(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, self.cfg.vocab_size), targets.view(-1))
        return logits, loss

if __name__ == "__main__":
    cfg = ModelConfig()
    m = DecoderLM(cfg)
    total_params = sum(p.numel() for p in m.parameters())
    print(f"Architecture loaded: {total_params / 1e6:.2f}M parameters")