# createsonline/ml/transformer.py
"""
CREATESONLINE Transformer Architecture — Pure NumPy Implementation

Full transformer encoder-decoder with:
  - Multi-head self-attention
  - Positional encoding (sinusoidal)
  - Feed-forward networks
  - Layer normalization
  - Residual connections
  - Causal (autoregressive) masking

This is a REAL implementation that can train small language models
on CPU with numpy only. Not suitable for large-scale models
(use PyTorch/JAX for that), but perfect for learning, prototyping,
and small-domain models.
"""

import math
import numpy as np
from typing import Optional, Tuple, List, Dict


# =====================================================
# LAYER NORMALIZATION
# =====================================================

class LayerNorm:
    """Layer normalization — stabilizes training"""
    
    def __init__(self, d_model: int, eps: float = 1e-6):
        self.d_model = d_model
        self.eps = eps
        self.gamma = np.ones(d_model)  # scale
        self.beta = np.zeros(d_model)  # shift
        # Gradients
        self.d_gamma = np.zeros_like(self.gamma)
        self.d_beta = np.zeros_like(self.beta)
        # Cache for backward
        self._cache = {}
    
    def forward(self, x: np.ndarray) -> np.ndarray:
        """x: (..., d_model)"""
        mean = x.mean(axis=-1, keepdims=True)
        var = x.var(axis=-1, keepdims=True)
        x_norm = (x - mean) / np.sqrt(var + self.eps)
        out = self.gamma * x_norm + self.beta
        self._cache = {'x_norm': x_norm, 'std': np.sqrt(var + self.eps)}
        return out
    
    def backward(self, d_out: np.ndarray) -> np.ndarray:
        """Compute gradients"""
        x_norm = self._cache['x_norm']
        std = self._cache['std']
        N = self.d_model
        
        self.d_gamma = np.sum(d_out * x_norm, axis=tuple(range(d_out.ndim - 1)))
        self.d_beta = np.sum(d_out, axis=tuple(range(d_out.ndim - 1)))
        
        dx_norm = d_out * self.gamma
        dx = (1.0 / N) / std * (
            N * dx_norm - dx_norm.sum(axis=-1, keepdims=True)
            - x_norm * (dx_norm * x_norm).sum(axis=-1, keepdims=True)
        )
        return dx
    
    def parameters(self) -> List[np.ndarray]:
        return [self.gamma, self.beta]
    
    def gradients(self) -> List[np.ndarray]:
        return [self.d_gamma, self.d_beta]


# =====================================================
# POSITIONAL ENCODING
# =====================================================

class PositionalEncoding:
    """
    Sinusoidal positional encoding (Vaswani et al., 2017).
    Adds position information to token embeddings.
    """
    
    def __init__(self, d_model: int, max_seq_len: int = 5000):
        self.d_model = d_model
        pe = np.zeros((max_seq_len, d_model))
        position = np.arange(0, max_seq_len).reshape(-1, 1)
        div_term = np.exp(np.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
        pe[:, 0::2] = np.sin(position * div_term)
        pe[:, 1::2] = np.cos(position * div_term[:d_model // 2]) if d_model % 2 else np.cos(position * div_term)
        self.pe = pe  # (max_seq_len, d_model)
    
    def forward(self, x: np.ndarray) -> np.ndarray:
        """x: (batch, seq_len, d_model)"""
        seq_len = x.shape[1]
        return x + self.pe[:seq_len]


# =====================================================
# LINEAR LAYER
# =====================================================

class Linear:
    """Dense / fully-connected layer"""
    
    def __init__(self, in_features: int, out_features: int, bias: bool = True):
        # Xavier initialization
        limit = math.sqrt(6.0 / (in_features + out_features))
        self.weight = np.random.uniform(-limit, limit, (in_features, out_features))
        self.bias = np.zeros(out_features) if bias else None
        self.has_bias = bias
        # Gradients
        self.d_weight = np.zeros_like(self.weight)
        self.d_bias = np.zeros_like(self.bias) if bias else None
        # Cache
        self._input = None
    
    def forward(self, x: np.ndarray) -> np.ndarray:
        """x: (..., in_features) -> (..., out_features)"""
        self._input = x
        out = x @ self.weight
        if self.has_bias:
            out = out + self.bias
        return out
    
    def backward(self, d_out: np.ndarray) -> np.ndarray:
        """Compute gradients and return d_input"""
        x = self._input
        # Reshape for batched matmul
        orig_shape = d_out.shape
        d_out_2d = d_out.reshape(-1, orig_shape[-1])
        x_2d = x.reshape(-1, x.shape[-1])
        
        self.d_weight = x_2d.T @ d_out_2d
        if self.has_bias:
            self.d_bias = d_out_2d.sum(axis=0)
        
        d_input = d_out @ self.weight.T
        return d_input
    
    def parameters(self) -> List[np.ndarray]:
        return [self.weight] + ([self.bias] if self.has_bias else [])
    
    def gradients(self) -> List[np.ndarray]:
        return [self.d_weight] + ([self.d_bias] if self.has_bias else [])


# =====================================================
# EMBEDDING LAYER
# =====================================================

class Embedding:
    """Token embedding lookup table"""
    
    def __init__(self, vocab_size: int, d_model: int):
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.weight = np.random.randn(vocab_size, d_model) * 0.02
        self.d_weight = np.zeros_like(self.weight)
        self._input_ids = None
    
    def forward(self, ids: np.ndarray) -> np.ndarray:
        """ids: (batch, seq_len) -> (batch, seq_len, d_model)"""
        self._input_ids = ids
        return self.weight[ids]
    
    def backward(self, d_out: np.ndarray) -> None:
        """Accumulate embedding gradients"""
        self.d_weight = np.zeros_like(self.weight)
        ids = self._input_ids
        # Scatter-add gradients
        np.add.at(self.d_weight, ids, d_out)
    
    def parameters(self) -> List[np.ndarray]:
        return [self.weight]
    
    def gradients(self) -> List[np.ndarray]:
        return [self.d_weight]


# =====================================================
# MULTI-HEAD SELF-ATTENTION
# =====================================================

class MultiHeadAttention:
    """
    Multi-head self-attention mechanism.
    
    Splits d_model into num_heads, computes scaled dot-product
    attention in parallel, then concatenates and projects.
    """
    
    def __init__(self, d_model: int, num_heads: int, dropout_rate: float = 0.0):
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.scale = math.sqrt(self.d_k)
        self.dropout_rate = dropout_rate
        
        # Q, K, V projections
        self.W_q = Linear(d_model, d_model)
        self.W_k = Linear(d_model, d_model)
        self.W_v = Linear(d_model, d_model)
        self.W_o = Linear(d_model, d_model)
        
        # Cache for backward
        self._cache = {}
    
    def _split_heads(self, x: np.ndarray) -> np.ndarray:
        """(batch, seq_len, d_model) -> (batch, num_heads, seq_len, d_k)"""
        batch, seq_len, _ = x.shape
        x = x.reshape(batch, seq_len, self.num_heads, self.d_k)
        return x.transpose(0, 2, 1, 3)
    
    def _merge_heads(self, x: np.ndarray) -> np.ndarray:
        """(batch, num_heads, seq_len, d_k) -> (batch, seq_len, d_model)"""
        batch, _, seq_len, _ = x.shape
        x = x.transpose(0, 2, 1, 3)
        return x.reshape(batch, seq_len, self.d_model)
    
    def forward(self, query: np.ndarray, key: np.ndarray, value: np.ndarray,
                mask: Optional[np.ndarray] = None) -> np.ndarray:
        """
        Compute multi-head attention.
        
        Args:
            query, key, value: (batch, seq_len, d_model)
            mask: Optional (batch, 1, 1, seq_len) or (1, 1, seq_len, seq_len)
        
        Returns:
            (batch, seq_len, d_model)
        """
        # Project Q, K, V
        Q = self._split_heads(self.W_q.forward(query))  # (B, H, S, dk)
        K = self._split_heads(self.W_k.forward(key))
        V = self._split_heads(self.W_v.forward(value))
        
        # Scaled dot-product attention
        scores = (Q @ K.transpose(0, 1, 3, 2)) / self.scale  # (B, H, S, S)
        
        if mask is not None:
            scores = scores + mask * (-1e9)
        
        attn_weights = self._softmax(scores)  # (B, H, S, S)
        
        # Dropout (training only)
        if self.dropout_rate > 0:
            drop_mask = (np.random.rand(*attn_weights.shape) > self.dropout_rate).astype(np.float64)
            attn_weights = attn_weights * drop_mask / (1 - self.dropout_rate + 1e-10)
        
        attn_output = attn_weights @ V  # (B, H, S, dk)
        
        # Merge heads and project
        concat = self._merge_heads(attn_output)  # (B, S, d_model)
        output = self.W_o.forward(concat)
        
        # Cache for backward
        self._cache = {
            'Q': Q, 'K': K, 'V': V,
            'attn_weights': attn_weights,
            'concat': concat,
            'query': query, 'key': key, 'value': value,
            'mask': mask,
        }
        
        return output
    
    def _softmax(self, x: np.ndarray) -> np.ndarray:
        """Numerically stable softmax along last axis"""
        e_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
        return e_x / (e_x.sum(axis=-1, keepdims=True) + 1e-10)
    
    def backward(self, d_out: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Backprop through multi-head attention"""
        Q, K, V = self._cache['Q'], self._cache['K'], self._cache['V']
        attn_weights = self._cache['attn_weights']
        
        # Through output projection
        d_concat = self.W_o.backward(d_out)
        
        # Reshape for heads
        batch = d_concat.shape[0]
        seq_len = d_concat.shape[1]
        d_attn_out = d_concat.reshape(batch, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3)
        
        # Through attention weights @ V
        d_attn_weights = d_attn_out @ V.transpose(0, 1, 3, 2)  # (B,H,S,S)
        d_V = attn_weights.transpose(0, 1, 3, 2) @ d_attn_out  # (B,H,S,dk)
        
        # Through softmax
        d_scores = attn_weights * (d_attn_weights - (d_attn_weights * attn_weights).sum(axis=-1, keepdims=True))
        d_scores = d_scores / self.scale
        
        # Through Q @ K^T
        d_Q = d_scores @ K       # (B,H,S,dk)
        d_K = d_scores.transpose(0, 1, 3, 2) @ Q  # (B,H,S,dk)
        
        # Merge heads back
        d_Q = self._merge_heads(d_Q)
        d_K = self._merge_heads(d_K)
        d_V = self._merge_heads(d_V)
        
        # Through input projections
        d_query = self.W_q.backward(d_Q)
        d_key = self.W_k.backward(d_K)
        d_value = self.W_v.backward(d_V)
        
        return d_query, d_key, d_value
    
    def parameters(self) -> List[np.ndarray]:
        params = []
        for layer in [self.W_q, self.W_k, self.W_v, self.W_o]:
            params.extend(layer.parameters())
        return params
    
    def gradients(self) -> List[np.ndarray]:
        grads = []
        for layer in [self.W_q, self.W_k, self.W_v, self.W_o]:
            grads.extend(layer.gradients())
        return grads


# =====================================================
# FEED-FORWARD NETWORK
# =====================================================

class FeedForward:
    """
    Position-wise feed-forward network.
    FFN(x) = ReLU(x·W1 + b1)·W2 + b2
    """
    
    def __init__(self, d_model: int, d_ff: int, dropout_rate: float = 0.0):
        self.linear1 = Linear(d_model, d_ff)
        self.linear2 = Linear(d_ff, d_model)
        self.dropout_rate = dropout_rate
        self._cache = {}
    
    def forward(self, x: np.ndarray) -> np.ndarray:
        hidden = self.linear1.forward(x)
        # GELU activation (smoother than ReLU, used in modern transformers)
        activated = self._gelu(hidden)
        if self.dropout_rate > 0:
            mask = (np.random.rand(*activated.shape) > self.dropout_rate).astype(np.float64)
            activated = activated * mask / (1 - self.dropout_rate + 1e-10)
        output = self.linear2.forward(activated)
        self._cache = {'hidden': hidden, 'activated': activated}
        return output
    
    def _gelu(self, x: np.ndarray) -> np.ndarray:
        """Gaussian Error Linear Unit"""
        return 0.5 * x * (1.0 + np.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * x**3)))
    
    def _gelu_backward(self, x: np.ndarray) -> np.ndarray:
        """Derivative of GELU"""
        cdf = 0.5 * (1.0 + np.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * x**3)))
        pdf = np.exp(-0.5 * x**2) / math.sqrt(2 * math.pi)
        return cdf + x * pdf
    
    def backward(self, d_out: np.ndarray) -> np.ndarray:
        hidden = self._cache['hidden']
        d_activated = self.linear2.backward(d_out)
        d_hidden = d_activated * self._gelu_backward(hidden)
        d_input = self.linear1.backward(d_hidden)
        return d_input
    
    def parameters(self) -> List[np.ndarray]:
        return self.linear1.parameters() + self.linear2.parameters()
    
    def gradients(self) -> List[np.ndarray]:
        return self.linear1.gradients() + self.linear2.gradients()


# =====================================================
# TRANSFORMER ENCODER BLOCK
# =====================================================

class TransformerEncoderBlock:
    """
    Single transformer encoder block:
        x -> LayerNorm -> MultiHeadAttention -> + residual
          -> LayerNorm -> FeedForward -> + residual
    """
    
    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout_rate: float = 0.0):
        self.attn = MultiHeadAttention(d_model, num_heads, dropout_rate)
        self.ff = FeedForward(d_model, d_ff, dropout_rate)
        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        self._cache = {}
    
    def forward(self, x: np.ndarray, mask: Optional[np.ndarray] = None) -> np.ndarray:
        # Self-attention with residual
        normed = self.norm1.forward(x)
        attn_out = self.attn.forward(normed, normed, normed, mask)
        x = x + attn_out
        self._cache['after_attn'] = x.copy()
        
        # Feed-forward with residual
        normed = self.norm2.forward(x)
        ff_out = self.ff.forward(normed)
        x = x + ff_out
        return x
    
    def backward(self, d_out: np.ndarray) -> np.ndarray:
        # Through FF residual
        d_ff = d_out.copy()
        d_normed2 = self.ff.backward(d_ff)
        d_x = d_out + self.norm2.backward(d_normed2)
        
        # Through attention residual
        d_attn = d_x.copy()
        d_normed1 = self.attn.backward(d_attn)[0]  # self-attention: q=k=v
        d_x = d_x + self.norm1.backward(d_normed1)
        
        return d_x
    
    def parameters(self) -> List[np.ndarray]:
        return (self.attn.parameters() + self.ff.parameters() +
                self.norm1.parameters() + self.norm2.parameters())
    
    def gradients(self) -> List[np.ndarray]:
        return (self.attn.gradients() + self.ff.gradients() +
                self.norm1.gradients() + self.norm2.gradients())


# =====================================================
# TRANSFORMER DECODER BLOCK
# =====================================================

class TransformerDecoderBlock:
    """
    Single transformer decoder block:
        x -> LayerNorm -> Masked Self-Attention -> + residual
          -> LayerNorm -> Cross-Attention (with encoder output) -> + residual
          -> LayerNorm -> FeedForward -> + residual
    """
    
    def __init__(self, d_model: int, num_heads: int, d_ff: int, dropout_rate: float = 0.0):
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout_rate)
        self.cross_attn = MultiHeadAttention(d_model, num_heads, dropout_rate)
        self.ff = FeedForward(d_model, d_ff, dropout_rate)
        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        self.norm3 = LayerNorm(d_model)
    
    def forward(self, x: np.ndarray, encoder_output: np.ndarray,
                src_mask: Optional[np.ndarray] = None,
                tgt_mask: Optional[np.ndarray] = None) -> np.ndarray:
        # Masked self-attention
        normed = self.norm1.forward(x)
        attn_out = self.self_attn.forward(normed, normed, normed, tgt_mask)
        x = x + attn_out
        
        # Cross-attention with encoder output
        normed = self.norm2.forward(x)
        cross_out = self.cross_attn.forward(normed, encoder_output, encoder_output, src_mask)
        x = x + cross_out
        
        # Feed-forward
        normed = self.norm3.forward(x)
        ff_out = self.ff.forward(normed)
        x = x + ff_out
        
        return x
    
    def backward(self, d_out: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """Returns (d_input, d_encoder_output)"""
        # FF
        d_normed3 = self.ff.backward(d_out)
        d_x = d_out + self.norm3.backward(d_normed3)
        
        # Cross-attention
        d_normed2 = d_x.copy()
        d_q, d_enc_k, d_enc_v = self.cross_attn.backward(d_normed2)
        d_x = d_x + self.norm2.backward(d_q)
        d_encoder = d_enc_k + d_enc_v
        
        # Self-attention
        d_normed1 = d_x.copy()
        d_self = self.self_attn.backward(d_normed1)[0]
        d_x = d_x + self.norm1.backward(d_self)
        
        return d_x, d_encoder
    
    def parameters(self) -> List[np.ndarray]:
        return (self.self_attn.parameters() + self.cross_attn.parameters() +
                self.ff.parameters() + self.norm1.parameters() +
                self.norm2.parameters() + self.norm3.parameters())
    
    def gradients(self) -> List[np.ndarray]:
        return (self.self_attn.gradients() + self.cross_attn.gradients() +
                self.ff.gradients() + self.norm1.gradients() +
                self.norm2.gradients() + self.norm3.gradients())


# =====================================================
# FULL TRANSFORMER MODEL
# =====================================================

class TransformerLM:
    """
    Transformer Language Model (decoder-only, like GPT).
    
    Can be trained for text generation on small datasets.
    
    Usage:
        model = TransformerLM(vocab_size=1000, d_model=128, num_heads=4, num_layers=2)
        logits = model.forward(input_ids)  # (batch, seq_len, vocab_size)
        loss = model.compute_loss(logits, target_ids)
        model.backward(logits, target_ids)
    """
    
    def __init__(self, vocab_size: int, d_model: int = 256, num_heads: int = 4,
                 num_layers: int = 4, d_ff: int = None, max_seq_len: int = 512,
                 dropout_rate: float = 0.1):
        
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.d_ff = d_ff or d_model * 4
        self.max_seq_len = max_seq_len
        self.dropout_rate = dropout_rate
        
        # Embedding + positional encoding
        self.token_embedding = Embedding(vocab_size, d_model)
        self.pos_encoding = PositionalEncoding(d_model, max_seq_len)
        
        # Decoder blocks
        self.layers = [
            TransformerEncoderBlock(d_model, num_heads, self.d_ff, dropout_rate)
            for _ in range(num_layers)
        ]
        
        # Final layer norm + output projection
        self.final_norm = LayerNorm(d_model)
        self.output_proj = Linear(d_model, vocab_size, bias=False)
        
        # Cache
        self._cache = {}
    
    def _causal_mask(self, seq_len: int) -> np.ndarray:
        """Create causal (autoregressive) attention mask"""
        # 1 where we should mask (upper triangle), 0 where we should attend
        mask = np.triu(np.ones((seq_len, seq_len)), k=1)
        return mask.reshape(1, 1, seq_len, seq_len)
    
    def forward(self, input_ids: np.ndarray) -> np.ndarray:
        """
        Forward pass.
        
        Args:
            input_ids: (batch, seq_len) integer token IDs
        
        Returns:
            logits: (batch, seq_len, vocab_size)
        """
        batch, seq_len = input_ids.shape
        
        # Token embeddings + positional encoding
        x = self.token_embedding.forward(input_ids)
        x = x * math.sqrt(self.d_model)  # Scale embeddings
        x = self.pos_encoding.forward(x)
        
        # Causal mask
        mask = self._causal_mask(seq_len)
        
        # Through decoder layers
        self._cache['layer_inputs'] = [x.copy()]
        for layer in self.layers:
            x = layer.forward(x, mask)
            self._cache['layer_inputs'].append(x.copy())
        
        # Final norm + projection to vocabulary
        x = self.final_norm.forward(x)
        logits = self.output_proj.forward(x)
        
        self._cache['final_normed'] = x
        self._cache['input_ids'] = input_ids
        
        return logits
    
    def compute_loss(self, logits: np.ndarray, targets: np.ndarray) -> float:
        """
        Compute cross-entropy loss for language modeling.
        
        Args:
            logits: (batch, seq_len, vocab_size)
            targets: (batch, seq_len) target token IDs
        
        Returns:
            scalar loss
        """
        batch, seq_len, vocab_size = logits.shape
        
        # Softmax
        probs = self._softmax(logits)
        self._cache['probs'] = probs
        
        # Cross-entropy loss
        # Clip for numerical stability
        probs_clipped = np.clip(probs, 1e-10, 1.0)
        
        # Gather target probabilities
        flat_probs = probs_clipped.reshape(-1, vocab_size)
        flat_targets = targets.reshape(-1)
        target_probs = flat_probs[np.arange(len(flat_targets)), flat_targets]
        
        loss = -np.mean(np.log(target_probs))
        return float(loss)
    
    def backward(self, logits: np.ndarray, targets: np.ndarray):
        """
        Backward pass through the entire model.
        Computes gradients for all parameters.
        """
        batch, seq_len, vocab_size = logits.shape
        probs = self._cache['probs']
        
        # Gradient of cross-entropy with softmax
        d_logits = probs.copy()
        flat_targets = targets.reshape(-1)
        d_logits_flat = d_logits.reshape(-1, vocab_size)
        d_logits_flat[np.arange(len(flat_targets)), flat_targets] -= 1.0
        d_logits = d_logits_flat.reshape(batch, seq_len, vocab_size) / (batch * seq_len)
        
        # Through output projection
        d_normed = self.output_proj.backward(d_logits)
        
        # Through final norm
        d_x = self.final_norm.backward(d_normed)
        
        # Through decoder layers (reverse order)
        for layer in reversed(self.layers):
            d_x = layer.backward(d_x)
        
        # Through positional encoding (no parameters)
        # Through scaling
        d_x = d_x * math.sqrt(self.d_model)
        
        # Through token embedding
        self.token_embedding.backward(d_x)
    
    def _softmax(self, x: np.ndarray) -> np.ndarray:
        """Numerically stable softmax"""
        e_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
        return e_x / (e_x.sum(axis=-1, keepdims=True) + 1e-10)
    
    def generate(self, prompt_ids: np.ndarray, max_new_tokens: int = 50,
                 temperature: float = 1.0, top_k: int = 0) -> np.ndarray:
        """
        Autoregressive text generation.
        
        Args:
            prompt_ids: (1, prompt_len) starting token IDs
            max_new_tokens: number of tokens to generate
            temperature: sampling temperature (lower = more deterministic)
            top_k: if > 0, only sample from top-k tokens
        
        Returns:
            (1, prompt_len + max_new_tokens) generated token IDs
        """
        ids = prompt_ids.copy()
        
        for _ in range(max_new_tokens):
            # Truncate to max sequence length
            input_ids = ids[:, -self.max_seq_len:]
            
            # Forward pass
            logits = self.forward(input_ids)
            
            # Get logits for last position
            next_logits = logits[:, -1, :] / max(temperature, 1e-8)
            
            # Top-k filtering
            if top_k > 0:
                top_k_idx = np.argsort(next_logits[0])[-top_k:]
                mask = np.full_like(next_logits, -1e10)
                mask[0, top_k_idx] = next_logits[0, top_k_idx]
                next_logits = mask
            
            # Sample
            probs = self._softmax(next_logits)
            next_id = np.array([[np.random.choice(self.vocab_size, p=probs[0])]])
            
            ids = np.concatenate([ids, next_id], axis=1)
        
        return ids
    
    def parameters(self) -> List[np.ndarray]:
        """Collect all parameters"""
        params = self.token_embedding.parameters()
        for layer in self.layers:
            params.extend(layer.parameters())
        params.extend(self.final_norm.parameters())
        params.extend(self.output_proj.parameters())
        return params
    
    def gradients(self) -> List[np.ndarray]:
        """Collect all gradients"""
        grads = self.token_embedding.gradients()
        for layer in self.layers:
            grads.extend(layer.gradients())
        grads.extend(self.final_norm.gradients())
        grads.extend(self.output_proj.gradients())
        return grads
    
    def num_parameters(self) -> int:
        """Count total number of trainable parameters"""
        return sum(p.size for p in self.parameters())
    
    def save(self, filepath: str):
        """Save model weights"""
        params = self.parameters()
        np.savez(filepath, *params,
                 config=np.array([self.vocab_size, self.d_model, self.num_heads,
                                  self.num_layers, self.d_ff, self.max_seq_len]))
    
    @classmethod
    def load(cls, filepath: str) -> 'TransformerLM':
        """Load model weights"""
        data = np.load(filepath + '.npz', allow_pickle=True)
        config = data['config']
        model = cls(
            vocab_size=int(config[0]),
            d_model=int(config[1]),
            num_heads=int(config[2]),
            num_layers=int(config[3]),
            d_ff=int(config[4]),
            max_seq_len=int(config[5])
        )
        params = model.parameters()
        for i, p in enumerate(params):
            p[:] = data[f'arr_{i}']
        return model
    
    def __repr__(self) -> str:
        return (f"TransformerLM(vocab={self.vocab_size}, d_model={self.d_model}, "
                f"heads={self.num_heads}, layers={self.num_layers}, "
                f"params={self.num_parameters():,})")


# =====================================================
# TRANSFORMER ENCODER-DECODER (SEQ2SEQ)
# =====================================================

class TransformerSeq2Seq:
    """
    Full encoder-decoder transformer for sequence-to-sequence tasks
    (translation, summarization, etc.)
    """
    
    def __init__(self, src_vocab_size: int, tgt_vocab_size: int,
                 d_model: int = 256, num_heads: int = 4, num_layers: int = 4,
                 d_ff: int = None, max_seq_len: int = 512, dropout_rate: float = 0.1):
        
        self.d_model = d_model
        self.src_vocab_size = src_vocab_size
        self.tgt_vocab_size = tgt_vocab_size
        self.d_ff = d_ff or d_model * 4
        
        # Encoder
        self.src_embedding = Embedding(src_vocab_size, d_model)
        self.src_pos = PositionalEncoding(d_model, max_seq_len)
        self.encoder_layers = [
            TransformerEncoderBlock(d_model, num_heads, self.d_ff, dropout_rate)
            for _ in range(num_layers)
        ]
        self.encoder_norm = LayerNorm(d_model)
        
        # Decoder
        self.tgt_embedding = Embedding(tgt_vocab_size, d_model)
        self.tgt_pos = PositionalEncoding(d_model, max_seq_len)
        self.decoder_layers = [
            TransformerDecoderBlock(d_model, num_heads, self.d_ff, dropout_rate)
            for _ in range(num_layers)
        ]
        self.decoder_norm = LayerNorm(d_model)
        
        # Output
        self.output_proj = Linear(d_model, tgt_vocab_size, bias=False)
    
    def encode(self, src_ids: np.ndarray, src_mask: Optional[np.ndarray] = None) -> np.ndarray:
        """Encode source sequence"""
        x = self.src_embedding.forward(src_ids) * math.sqrt(self.d_model)
        x = self.src_pos.forward(x)
        for layer in self.encoder_layers:
            x = layer.forward(x, src_mask)
        return self.encoder_norm.forward(x)
    
    def decode(self, tgt_ids: np.ndarray, encoder_output: np.ndarray,
               src_mask: Optional[np.ndarray] = None,
               tgt_mask: Optional[np.ndarray] = None) -> np.ndarray:
        """Decode target sequence"""
        seq_len = tgt_ids.shape[1]
        if tgt_mask is None:
            tgt_mask = np.triu(np.ones((seq_len, seq_len)), k=1).reshape(1, 1, seq_len, seq_len)
        
        x = self.tgt_embedding.forward(tgt_ids) * math.sqrt(self.d_model)
        x = self.tgt_pos.forward(x)
        for layer in self.decoder_layers:
            x = layer.forward(x, encoder_output, src_mask, tgt_mask)
        x = self.decoder_norm.forward(x)
        return self.output_proj.forward(x)
    
    def forward(self, src_ids: np.ndarray, tgt_ids: np.ndarray) -> np.ndarray:
        """Full forward pass"""
        encoder_output = self.encode(src_ids)
        logits = self.decode(tgt_ids, encoder_output)
        return logits
    
    def parameters(self) -> List[np.ndarray]:
        params = self.src_embedding.parameters() + self.tgt_embedding.parameters()
        for layer in self.encoder_layers:
            params.extend(layer.parameters())
        for layer in self.decoder_layers:
            params.extend(layer.parameters())
        params.extend(self.encoder_norm.parameters() + self.decoder_norm.parameters())
        params.extend(self.output_proj.parameters())
        return params
    
    def gradients(self) -> List[np.ndarray]:
        grads = self.src_embedding.gradients() + self.tgt_embedding.gradients()
        for layer in self.encoder_layers:
            grads.extend(layer.gradients())
        for layer in self.decoder_layers:
            grads.extend(layer.gradients())
        grads.extend(self.encoder_norm.gradients() + self.decoder_norm.gradients())
        grads.extend(self.output_proj.gradients())
        return grads
    
    def num_parameters(self) -> int:
        return sum(p.size for p in self.parameters())
