import numpy as np


def softmax(x, axis=-1):
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e_x / np.sum(e_x, axis=axis, keepdims=True)


def relu(x):
    return np.maximum(0, x)


class Linear:
    def __init__(self, in_features, out_features):
        self.weight = np.random.randn(in_features, out_features) * np.sqrt(2.0 / in_features)
        self.bias = np.zeros(out_features)
        
    def forward(self, x):
        return x @ self.weight + self.bias


class LayerNorm:
    def __init__(self, d_model, eps=1e-6):
        self.gamma = np.ones(d_model)
        self.beta = np.zeros(d_model)
        self.eps = eps
        
    def forward(self, x):
        mean = np.mean(x, axis=-1, keepdims=True)
        std = np.std(x, axis=-1, keepdims=True)
        return self.gamma * (x - mean) / (std + self.eps) + self.beta


class MultiHeadAttention:
    def __init__(self, d_model, num_heads):
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        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)
        
    def forward(self, Q, K, V, mask=None):
        batch_size = Q.shape[0]
        seq_len = Q.shape[1]
        
        Q = self.W_q.forward(Q).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3)
        K = self.W_k.forward(K).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3)
        V = self.W_v.forward(V).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(0, 2, 1, 3)
        
        attn_scores = Q @ K.transpose(0, 1, 3, 2) / np.sqrt(self.d_k)
        if mask is not None:
            attn_scores = np.where(mask == 0, -1e9, attn_scores)
        attn_probs = softmax(attn_scores, axis=-1)
        attn_output = attn_probs @ V
        
        attn_output = attn_output.transpose(0, 2, 1, 3).reshape(batch_size, seq_len, self.d_model)
        output = self.W_o.forward(attn_output)
        return output


class FeedForward:
    def __init__(self, d_model, d_ff):
        self.linear1 = Linear(d_model, d_ff)
        self.linear2 = Linear(d_ff, d_model)
        
    def forward(self, x):
        return self.linear2.forward(relu(self.linear1.forward(x)))


class TransformerBlock:
    def __init__(self, d_model, num_heads, d_ff):
        self.attention = MultiHeadAttention(d_model, num_heads)
        self.feed_forward = FeedForward(d_model, d_ff)
        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        
    def forward(self, x, mask=None):
        attn_output = self.attention.forward(x, x, x, mask)
        x = self.norm1.forward(x + attn_output)
        ff_output = self.feed_forward.forward(x)
        x = self.norm2.forward(x + ff_output)
        return x


class SimpleTransformer:
    def __init__(self, vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len):
        self.d_model = d_model
        self.embedding = np.random.randn(vocab_size, d_model) * 0.01
        self.positional_encoding = self._create_positional_encoding(max_seq_len, d_model)
        self.transformer_blocks = [
            TransformerBlock(d_model, num_heads, d_ff) for _ in range(num_layers)
        ]
        
    def _create_positional_encoding(self, max_seq_len, d_model):
        pe = np.zeros((max_seq_len, d_model))
        position = np.arange(0, max_seq_len).reshape(-1, 1).astype(float)
        div_term = np.exp(np.arange(0, d_model, 2).astype(float) * -(np.log(10000.0) / d_model))
        pe[:, 0::2] = np.sin(position * div_term)
        pe[:, 1::2] = np.cos(position * div_term)
        return pe
        
    def forward(self, x, mask=None):
        seq_len = x.shape[1]
        x = self.embedding[x] * np.sqrt(self.d_model)
        x = x + self.positional_encoding[:seq_len, :]
        
        for block in self.transformer_blocks:
            x = block.forward(x, mask)
            
        return x
    
    def count_parameters(self):
        count = 0
        count += self.embedding.size
        for block in self.transformer_blocks:
            count += block.attention.W_q.weight.size + block.attention.W_q.bias.size
            count += block.attention.W_k.weight.size + block.attention.W_k.bias.size
            count += block.attention.W_v.weight.size + block.attention.W_v.bias.size
            count += block.attention.W_o.weight.size + block.attention.W_o.bias.size
            count += block.feed_forward.linear1.weight.size + block.feed_forward.linear1.bias.size
            count += block.feed_forward.linear2.weight.size + block.feed_forward.linear2.bias.size
            count += block.norm1.gamma.size + block.norm1.beta.size
            count += block.norm2.gamma.size + block.norm2.beta.size
        return count


def create_padding_mask(seq, pad_idx=0):
    return (seq != pad_idx).astype(float)[:, np.newaxis, np.newaxis, :]


def create_look_ahead_mask(size):
    mask = np.triu(np.ones((size, size)), k=1).astype(bool)
    return (~mask).astype(float)


if __name__ == "__main__":
    vocab_size = 1000
    d_model = 512
    num_heads = 8
    num_layers = 6
    d_ff = 2048
    max_seq_len = 100
    
    model = SimpleTransformer(vocab_size, d_model, num_heads, num_layers, d_ff, max_seq_len)
    
    batch_size = 2
    seq_len = 10
    x = np.random.randint(0, vocab_size, (batch_size, seq_len))
    
    output = model.forward(x)
    print(f"Input shape: {x.shape}")
    print(f"Output shape: {output.shape}")
    print(f"Model parameters: {model.count_parameters():,}")