import numpy as np
from .utils import silu, softmax, rms_norm

def precompute_rope_freqs(dim, max_seq_len, theta=10000.0):
    """Precompute rotary positional embedding frequencies."""
    freqs = 1.0 / (theta ** (np.arange(0, dim, 2)[: (dim // 2)] / dim))
    t = np.arange(max_seq_len)
    freqs = np.outer(t, freqs) # (max_seq_len, dim/2)
    return np.cos(freqs), np.sin(freqs)

def apply_rope(x, cos, sin):
    """Apply rotary positional embeddings to a tensor."""
    # x shape: (B, nh, T, hs)
    # cos/sin shape: (T, hs/2)
    T = x.shape[2]
    cos = cos[:T].reshape(1, 1, T, -1) # (1, 1, T, hs/2)
    sin = sin[:T].reshape(1, 1, T, -1) # (1, 1, T, hs/2)
    
    # Reshape x to apply rotation on pairs
    x_real = x[..., 0::2]
    x_imag = x[..., 1::2]
    
    # Rotate
    # (a + ib)(cos + isin) = (a*cos - b*sin) + i(a*sin + b*cos)
    x_out_real = x_real * cos - x_imag * sin
    x_out_imag = x_real * sin + x_imag * cos
    
    # Stack back
    x_out = np.empty_like(x)
    x_out[..., 0::2] = x_out_real
    x_out[..., 1::2] = x_out_imag
    return x_out

class Linear:
    def __init__(self, in_features, out_features, bias=True):
        self.w = np.random.randn(in_features, out_features) * 0.02
        self.b = np.zeros(out_features) if bias else None
        
    def __call__(self, x):
        res = x @ self.w
        if self.b is not None:
            res += self.b
        return res

class RMSNorm:
    def __init__(self, dim, eps=1e-6):
        self.eps = eps
        self.weight = np.ones(dim)
        
    def __call__(self, x):
        return rms_norm(x, self.weight, self.eps)

class MultiHeadAttention:
    def __init__(self, n_embd, n_head, bias=False):
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.n_embd = n_embd
        self.head_dim = n_embd // n_head
        
        self.wq = Linear(n_embd, n_embd, bias=bias)
        self.wk = Linear(n_embd, n_embd, bias=bias)
        self.wv = Linear(n_embd, n_embd, bias=bias)
        self.wo = Linear(n_embd, n_embd, bias=bias)
        
    def __call__(self, x, rope_cos, rope_sin):
        B, T, C = x.shape
        
        xq = self.wq(x).reshape(B, T, self.n_head, self.head_dim).transpose(0, 2, 1, 3)
        xk = self.wk(x).reshape(B, T, self.n_head, self.head_dim).transpose(0, 2, 1, 3)
        xv = self.wv(x).reshape(B, T, self.n_head, self.head_dim).transpose(0, 2, 1, 3)
        
        # Apply RoPE
        xq = apply_rope(xq, rope_cos, rope_sin)
        xk = apply_rope(xk, rope_cos, rope_sin)
        
        # Causal mask
        mask = np.tril(np.ones((T, T))).reshape(1, 1, T, T)
        
        # Scaled dot-product attention
        scores = (xq @ xk.transpose(0, 1, 3, 2)) / np.sqrt(self.head_dim)
        scores = np.where(mask == 0, -1e10, scores)
        probs = softmax(scores, axis=-1)
        
        output = probs @ xv # (B, nh, T, head_dim)
        output = output.transpose(0, 2, 1, 3).reshape(B, T, C)
        
        return self.wo(output)

class SwigluMLP:
    def __init__(self, n_embd, intermediate_size=None, bias=False):
        if intermediate_size is None:
            intermediate_size = int(8/3 * n_embd)
            # Round to nearest multiple of 256 for efficiency (Llama style)
            intermediate_size = 256 * ((intermediate_size + 256 - 1) // 256)
            
        self.w1 = Linear(n_embd, intermediate_size, bias=bias)
        self.w2 = Linear(intermediate_size, n_embd, bias=bias)
        self.w3 = Linear(n_embd, intermediate_size, bias=bias)
        
    def __call__(self, x):
        # Swiglu: (silu(w1(x)) * w3(x)) @ w2
        return self.w2(silu(self.w1(x)) * self.w3(x))

class Block:
    def __init__(self, n_embd, n_head, bias=False):
        self.attention_norm = RMSNorm(n_embd)
        self.attention = MultiHeadAttention(n_embd, n_head, bias=bias)
        self.ffn_norm = RMSNorm(n_embd)
        self.feed_forward = SwigluMLP(n_embd, bias=bias)
        
    def __call__(self, x, rope_cos, rope_sin):
        h = x + self.attention(self.attention_norm(x), rope_cos, rope_sin)
        out = h + self.feed_forward(self.ffn_norm(h))
        return out

class PotatoGPT:
    def __init__(self, vocab_size, n_layer=4, n_head=4, n_embd=128, block_size=128, bias=False):
        self.block_size = block_size
        self.tok_embeddings = np.random.randn(vocab_size, n_embd) * 0.02
        self.layers = [Block(n_embd, n_head, bias=bias) for _ in range(n_layer)]
        self.norm = RMSNorm(n_embd)
        self.output = Linear(n_embd, vocab_size, bias=False)
        
        # Precompute RoPE frequencies
        self.rope_cos, self.rope_sin = precompute_rope_freqs(n_embd // n_head, block_size)
        
    def __call__(self, idx):
        B, T = idx.shape
        h = self.tok_embeddings[idx] # (B, T, n_embd)
        
        # Get RoPE for current T
        cos = self.rope_cos[:T]
        sin = self.rope_sin[:T]
        
        for layer in self.layers:
            h = layer(h, cos, sin)
            
        h = self.norm(h)
        logits = self.output(h)
        return logits
