"""
MicroGLM: Tiny GLM-style model in PyTorch.
GLM (General Language Model) uses a prefix-LM architecture:
bidirectional attention on a prefix span + autoregressive generation on the rest.
This is a minimal implementation that fits in 6GB VRAM.

Reference: "GLM: General Language Model Pretraining with Autoregressive Blank Infilling" (Du et al., 2021)
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F


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

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


class GLMAttention(nn.Module):
    """
    Attention with support for 2D attention masks (for prefix-LM).
    The mask has shape (B, 1, T, T) where:
    - 0 means attend (bidirectional / causal allowed)
    - -inf means blocked
    """
    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.head_dim = n_embd // n_head
        self.qkv = nn.Linear(n_embd, n_embd * 3, bias=False)
        self.proj = nn.Linear(n_embd, n_embd, bias=False)
        self.attn_drop = nn.Dropout(dropout)
        self.resid_drop = nn.Dropout(dropout)

    def forward(self, x, attn_mask=None):
        """
        x: (B, T, C)
        attn_mask: (B, 1, T, T) or None (uses default causal mask)
        """
        B, T, C = x.shape
        qkv = self.qkv(x).reshape(B, T, 3, self.n_head, self.head_dim).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)

        if attn_mask is not None:
            # attn_mask: 0 = attend, -inf = block
            att = att + attn_mask
        else:
            # Default causal mask
            mask = torch.tril(torch.ones(T, T, device=x.device)).view(1, 1, T, T)
            att = att.masked_fill(mask == 0, float('-inf'))

        att = F.softmax(att, dim=-1)
        att = self.attn_drop(att)
        y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_drop(self.proj(y))


class GLMMLP(nn.Module):
    def __init__(self, n_embd, dropout):
        super().__init__()
        self.fc = nn.Linear(n_embd, 4 * n_embd, bias=False)
        self.proj = nn.Linear(4 * n_embd, n_embd, bias=False)
        self.drop = nn.Dropout(dropout)

    def forward(self, x):
        x = F.gelu(self.fc(x))
        x = self.proj(x)
        return self.drop(x)


class GLMBlock(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        self.ln1 = RMSNorm(n_embd)
        self.attn = GLMAttention(n_embd, n_head, block_size, dropout)
        self.ln2 = RMSNorm(n_embd)
        self.mlp = GLMMLP(n_embd, dropout)

    def forward(self, x, attn_mask=None):
        x = x + self.attn(self.ln1(x), attn_mask)
        x = x + self.mlp(self.ln2(x))
        return x


class MicroGLM(nn.Module):
    """
    Micro GLM: Prefix-LM decoder.
    Supports two attention modes:
    - Causal LM (default): standard autoregressive generation
    - Prefix LM: bidirectional attention on first `prefix_len` tokens, causal on the rest

    This is controlled by passing a custom attention mask during forward().
    For training, we create a 2D mask where the prefix region is bidirectional
    and the suffix region is causal.
    """

    def __init__(self, vocab_size, block_size, n_layer=2, n_head=4, n_embd=128, dropout=0.1):
        super().__init__()
        self.block_size = block_size
        self.wte = nn.Embedding(vocab_size, n_embd)
        self.wpe = nn.Embedding(block_size, n_embd)
        self.blocks = nn.ModuleList([GLMBlock(n_embd, n_head, block_size, dropout) for _ in range(n_layer)])
        self.ln_f = RMSNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
        self.lm_head.weight = self.wte.weight
        self.apply(self._init_weights)

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

    def _build_prefix_lm_mask(self, T, prefix_len, device):
        """
        Build a 2D attention mask for prefix-LM.
        - For positions i < prefix_len and j < prefix_len: bidirectional (0)
        - For positions i >= prefix_len: causal (attend only to j <= i)
        - All other positions: -inf

        Returns: (1, 1, T, T) mask where 0 = allowed, -inf = blocked
        """
        # Start with causal mask
        mask = torch.tril(torch.ones(T, T, device=device))
        # Set the prefix block to be fully connected (bidirectional)
        mask[:, :prefix_len] = 1.0
        # Convert to float mask: 0 = attend, -inf = blocked
        mask = mask.view(1, 1, T, T)
        mask = mask.masked_fill(mask == 0, float('-inf'))
        mask = mask.masked_fill(mask == 1.0, 0.0)
        return mask

    def forward(self, idx, targets=None, prefix_len=0):
        """
        idx: (B, T) input tokens
        targets: (B, T) target tokens (shifted for loss)
        prefix_len: number of prefix tokens to use bidirectional attention
        """
        B, T = idx.shape
        if T > self.block_size:
            raise ValueError(f'block size exceeded: {T} > {self.block_size}')
        pos = torch.arange(T, device=idx.device)
        x = self.wte(idx) + self.wpe(pos)

        if prefix_len > 0:
            attn_mask = self._build_prefix_lm_mask(T, prefix_len, idx.device)
        else:
            attn_mask = None

        for block in self.blocks:
            x = block(x, attn_mask)

        x = self.ln_f(x)
        logits = self.lm_head(x)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=40):
        """Standard causal generation (no prefix)."""
        self.eval()
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.block_size:]
            logits, _ = self(idx_cond, prefix_len=0)
            logits = logits[:, -1, :] / temperature
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('inf')
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, idx_next], dim=1)
        return idx

    @torch.no_grad()
    def generate_with_prefix(self, idx, prefix_len, max_new_tokens, temperature=1.0, top_k=40):
        """
        Generate tokens where the first `prefix_len` tokens use bidirectional attention
        and the rest are autoregressive.
        """
        self.eval()
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.block_size:]
            T = idx_cond.shape[1]
            prefix = min(prefix_len, T)
            logits, _ = self(idx_cond, prefix_len=prefix)
            logits = logits[:, -1, :] / temperature
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('inf')
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, idx_next], dim=1)
        return idx