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


@dataclass
class GLMConfig:
    vocab_size: int = 151552
    hidden_size: int = 4096
    intermediate_size: int = 14336
    num_layers: int = 32
    num_heads: int = 32
    num_kv_heads: int = 8      # GQA
    head_dim: int = 128
    eps: float = 1e-6
    rope_base: float = 10000.0


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

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


# ──────────────────────────────── RoPE ───────────────────────────────────
class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim, base=10000.0):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def forward(self, seq_len, device, dtype):
        t = torch.arange(seq_len, device=device).float()
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        # duplicate to match head_dim
        emb = torch.cat([freqs, freqs], dim=-1)
        return emb.cos().to(dtype), emb.sin().to(dtype)


def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2:]
    return torch.cat([-x2, x1], dim=-1)


def apply_rotary(q, k, cos, sin):
    cos = cos[None, None, :, :]
    sin = sin[None, None, :, :]
    q = (q * cos) + (rotate_half(q) * sin)
    k = (k * cos) + (rotate_half(k) * sin)
    return q, k


# ─────────────────────────── Self-Attention (GQA) ────────────────────────
class SelfAttention(nn.Module):
    def __init__(self, cfg: GLMConfig):
        super().__init__()
        self.cfg = cfg
        self.q_proj = nn.Linear(cfg.hidden_size, cfg.num_heads * cfg.head_dim, bias=False)
        self.k_proj = nn.Linear(cfg.hidden_size, cfg.num_kv_heads * cfg.head_dim, bias=False)
        self.v_proj = nn.Linear(cfg.hidden_size, cfg.num_kv_heads * cfg.head_dim, bias=False)
        self.o_proj = nn.Linear(cfg.num_heads * cfg.head_dim, cfg.hidden_size, bias=False)

    def forward(self, x, cos, sin):
        B, T, _ = x.shape
        q = self.q_proj(x).view(B, T, self.cfg.num_heads, self.cfg.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.cfg.num_kv_heads, self.cfg.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.cfg.num_kv_heads, self.cfg.head_dim).transpose(1, 2)

        q, k = apply_rotary(q, k, cos, sin)

        # GQA: replicate KV heads
        if self.cfg.num_kv_heads != self.cfg.num_heads:
            rep = self.cfg.num_heads // self.cfg.num_kv_heads
            k = k.repeat_interleave(rep, dim=1)
            v = v.repeat_interleave(rep, dim=1)

        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        out = out.transpose(1, 2).contiguous().view(B, T, -1)
        return self.o_proj(out)


# ───────────────────────────── MLP (SwiGLU) ──────────────────────────────
class MLP(nn.Module):
    def __init__(self, cfg: GLMConfig):
        super().__init__()
        self.gate_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
        self.up_proj   = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
        self.down_proj = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False)

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


# ───────────────────────────── Transformer Block ─────────────────────────
class GLMBlock(nn.Module):
    def __init__(self, cfg: GLMConfig):
        super().__init__()
        self.input_norm = RMSNorm(cfg.hidden_size, cfg.eps)
        self.self_attn  = SelfAttention(cfg)
        self.post_norm  = RMSNorm(cfg.hidden_size, cfg.eps)
        self.mlp        = MLP(cfg)

    def forward(self, x, cos, sin):
        # Pre-norm Attention
        x = x + self.self_attn(self.input_norm(x), cos, sin)
        # Pre-norm MLP
        x = x + self.mlp(self.post_norm(x))
        return x


# ─────────────────────────────── Full Model ──────────────────────────────
class GLMForCausalLM(nn.Module):
    def __init__(self, cfg: GLMConfig):
        super().__init__()
        self.cfg = cfg
        self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
        self.rotary_emb   = RotaryEmbedding(cfg.head_dim, cfg.rope_base)
        self.layers       = nn.ModuleList([GLMBlock(cfg) for _ in range(cfg.num_layers)])
        self.final_norm   = RMSNorm(cfg.hidden_size, cfg.eps)
        self.lm_head      = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
        # weight tying
        self.lm_head.weight = self.embed_tokens.weight

    def forward(self, input_ids):
        B, T = input_ids.shape
        x = self.embed_tokens(input_ids)
        cos, sin = self.rotary_emb(T, x.device, x.dtype)

        for layer in self.layers:
            x = layer(x, cos, sin)

        x = self.final_norm(x)
        return self.lm_head(x)


# ─────────────────────────────── Test run ────────────────────────────────
if __name__ == "__main__":
    cfg = GLMConfig()
    model = GLMForCausalLM(cfg)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"Total parameters: {n_params/1e9:.2f}B")

    ids = torch.randint(0, cfg.vocab_size, (2, 16))
    with torch.no_grad():
        logits = model(ids)
    print(f"Output logits shape: {logits.shape}")  # (2, 16, vocab_size)