"""Quark-50M model definition — standalone."""

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

@dataclass
class ModelConfig:
    vocab_size: int = 16384; d_model: int = 512; n_heads: int = 8
    n_kv_heads: int = 4; n_layers: int = 12; d_ff: int = 1408
    head_dim: int = 64; max_seq_len: int = 2048; rope_theta: float = 10000.0
    rms_eps: float = 1e-5; qkv_bias: bool = False; dropout: float = 0.0

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__(); self.eps = eps; self.scale = nn.Parameter(torch.ones(dim))
    def forward(self, x):
        return (x.float() * x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()).to(x.dtype) * self.scale

class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim, max_seq_len, theta=10000.0):
        super().__init__()
        inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False); self._build(max_seq_len)
    def _build(self, seq_len):
        t = torch.arange(seq_len, device=self.inv_freq.device).float()
        freqs = torch.outer(t, self.inv_freq); emb = torch.cat([freqs, freqs], dim=-1)
        self.register_buffer("cos_cache", emb.cos()[None, None], persistent=False)
        self.register_buffer("sin_cache", emb.sin()[None, None], persistent=False); self._max = seq_len
    @staticmethod
    def _rot(x):
        x1, x2 = x.chunk(2, dim=-1); return torch.cat([-x2, x1], dim=-1)
    def forward(self, q, k):
        T = q.size(2)
        if T > self._max: self._build(T)
        c, s = self.cos_cache[:,:,:T], self.sin_cache[:,:,:T]
        return q*c + self._rot(q)*s, k*c + self._rot(k)*s

class GQA(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.n_heads, self.n_kv_heads = cfg.n_heads, cfg.n_kv_heads
        self.n_groups, self.head_dim = cfg.n_heads // cfg.n_kv_heads, cfg.head_dim
        self.q_proj = nn.Linear(cfg.d_model, cfg.n_heads * cfg.head_dim, bias=cfg.qkv_bias)
        self.k_proj = nn.Linear(cfg.d_model, cfg.n_kv_heads * cfg.head_dim, bias=cfg.qkv_bias)
        self.v_proj = nn.Linear(cfg.d_model, cfg.n_kv_heads * cfg.head_dim, bias=cfg.qkv_bias)
        self.o_proj = nn.Linear(cfg.n_heads * cfg.head_dim, cfg.d_model, bias=False)
        self.rope = RotaryEmbedding(cfg.head_dim, cfg.max_seq_len, cfg.rope_theta)
    def forward(self, x):
        B, T, _ = x.shape
        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        q, k = self.rope(q, k)
        if self.n_groups > 1:
            B_r, _, T_r, D_r = k.shape
            k = k[:,:,None,:,:].expand(B_r, self.n_kv_heads, self.n_groups, T_r, D_r).reshape(B_r, self.n_heads, T_r, D_r)
            v = v[:,:,None,:,:].expand(B_r, self.n_kv_heads, self.n_groups, T_r, D_r).reshape(B_r, self.n_heads, T_r, D_r)
        return self.o_proj(F.scaled_dot_product_attention(q, k, v, is_causal=True).transpose(1, 2).contiguous().view(B, T, -1))

class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.norm_attn = RMSNorm(cfg.d_model, cfg.rms_eps); self.attn = GQA(cfg)
        self.norm_ffn = RMSNorm(cfg.d_model, cfg.rms_eps)
        self.gate = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
        self.up = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
        self.down = nn.Linear(cfg.d_ff, cfg.d_model, bias=False)
    def forward(self, x):
        x = x + self.attn(self.norm_attn(x))
        h = self.norm_ffn(x); x = x + self.down(F.silu(self.gate(h)) * self.up(h))
        return x

class Quark(nn.Module):
    def __init__(self, cfg):
        super().__init__(); self.cfg = cfg
        self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.layers = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layers)])
        self.norm = RMSNorm(cfg.d_model, cfg.rms_eps)
        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        self.lm_head.weight = self.embed_tokens.weight
    def forward(self, ids, labels=None):
        x = self.embed_tokens(ids)
        for layer in self.layers: x = layer(x)
        logits = self.lm_head(self.norm(x))
        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits.view(-1, self.cfg.vocab_size), labels.view(-1), ignore_index=-100)
        return loss, logits