"""
TinyLLM Model Architecture

A small transformer language model (~54.93M parameters) trained from scratch.
Architecture:
- 12 layers
- 512 hidden size
- 8 attention heads
- 1408 intermediate (FFN)
- 32000 vocab size
- 512 max sequence length
- RoPE position encoding
- RMSNorm
- SwiGLU activation
- Weight tying
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Dict, Any, Optional


# Model configuration
MODEL_CONFIG = {
    "vocab_size": 32000,
    "hidden_size": 512,
    "num_layers": 12,
    "num_heads": 8,
    "intermediate_size": 1408,
    "max_position_embeddings": 512,
    "dropout": 0.0,
    "tie_weights": True,
}


class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization."""
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
    
    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight


class RotaryEmbedding(nn.Module):
    """Rotary Position Embedding (RoPE)."""
    def __init__(self, dim: int, max_seq_len: int = 512, base: int = 10000):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        self.max_seq_len = max_seq_len
    
    def forward(self, seq_len: int, device):
        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        return emb.cos(), emb.sin()


def rotate_half(x):
    """Rotate half the hidden dims of the input."""
    x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q, k, cos, sin):
    """Apply rotary positional embeddings to query and key tensors."""
    cos = cos.unsqueeze(0).unsqueeze(0)
    sin = sin.unsqueeze(0).unsqueeze(0)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


class Attention(nn.Module):
    """Multi-head attention with RoPE."""
    def __init__(self, config: Dict[str, Any]):
        super().__init__()
        self.hidden_size = config["hidden_size"]
        self.num_heads = config["num_heads"]
        self.head_dim = self.hidden_size // self.num_heads
        
        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        self.k_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        
        self.rotary = RotaryEmbedding(self.head_dim, config["max_position_embeddings"])
        self.dropout = nn.Dropout(config.get("dropout", 0.0))
    
    def forward(self, x, attention_mask=None):
        B, T, C = x.shape
        
        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        
        cos, sin = self.rotary(T, x.device)
        q, k = apply_rotary_pos_emb(q, k, cos, sin)
        
        # Scaled dot-product attention
        attn_weights = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        
        # Causal mask
        causal_mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)
        attn_weights.masked_fill_(causal_mask, float('-inf'))
        
        if attention_mask is not None:
            attn_weights = attn_weights + attention_mask
        
        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(x.dtype)
        attn_weights = self.dropout(attn_weights)
        
        out = torch.matmul(attn_weights, v)
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.o_proj(out)


class FFN(nn.Module):
    """Feed-forward network with SwiGLU activation."""
    def __init__(self, config: Dict[str, Any]):
        super().__init__()
        # SwiGLU: w1=gate, w2=down, w3=up
        self.w1 = nn.Linear(config["hidden_size"], config["intermediate_size"], bias=False)
        self.w2 = nn.Linear(config["intermediate_size"], config["hidden_size"], bias=False)
        self.w3 = nn.Linear(config["hidden_size"], config["intermediate_size"], bias=False)
    
    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class TransformerBlock(nn.Module):
    """Transformer block with pre-norm architecture."""
    def __init__(self, config: Dict[str, Any]):
        super().__init__()
        self.norm1 = RMSNorm(config["hidden_size"])
        self.attn = Attention(config)
        self.norm2 = RMSNorm(config["hidden_size"])
        self.ffn = FFN(config)
    
    def forward(self, x, attention_mask=None):
        x = x + self.attn(self.norm1(x), attention_mask)
        x = x + self.ffn(self.norm2(x))
        return x


class TinyLLM(nn.Module):
    """
    TinyLLM: A small decoder-only transformer language model.
    
    Args:
        config: Dictionary containing model configuration
    """
    def __init__(self, config: Dict[str, Any] = None):
        super().__init__()
        self.config = config or MODEL_CONFIG
        
        self.embed_tokens = nn.Embedding(self.config["vocab_size"], self.config["hidden_size"])
        self.layers = nn.ModuleList([
            TransformerBlock(self.config) 
            for _ in range(self.config["num_layers"])
        ])
        self.norm = RMSNorm(self.config["hidden_size"])
        self.lm_head = nn.Linear(self.config["hidden_size"], self.config["vocab_size"], bias=False)
        
        # Tie embeddings if configured
        if self.config.get("tie_weights", True):
            self.lm_head.weight = self.embed_tokens.weight
        
        # Register causal mask buffer
        max_len = self.config["max_position_embeddings"]
        self.register_buffer("causal_mask", 
            torch.triu(torch.ones(max_len, max_len, dtype=torch.bool), diagonal=1))
    
    def forward(self, input_ids, attention_mask=None, labels=None):
        x = self.embed_tokens(input_ids)
        
        for layer in self.layers:
            x = layer(x, attention_mask)
        
        x = self.norm(x)
        logits = self.lm_head(x)
        
        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, self.config["vocab_size"]),
                shift_labels.view(-1),
                ignore_index=-100
            )
        
        return {"logits": logits, "loss": loss}
    
    @torch.no_grad()
    def generate(
        self, 
        input_ids: torch.Tensor,
        max_new_tokens: int = 100,
        temperature: float = 0.8,
        top_p: float = 0.9,
        top_k: int = 50,
        eos_token_id: Optional[int] = None,
        repetition_penalty: float = 1.0,
    ) -> torch.Tensor:
        """
        Generate text autoregressively.
        
        Args:
            input_ids: Input token IDs [batch_size, seq_len]
            max_new_tokens: Maximum number of tokens to generate
            temperature: Sampling temperature (higher = more random)
            top_p: Nucleus sampling threshold
            top_k: Top-k sampling threshold
            eos_token_id: Token ID that signals end of generation
            repetition_penalty: Penalty for repeating tokens
            
        Returns:
            Generated token IDs including the prompt
        """
        self.eval()
        
        for _ in range(max_new_tokens):
            # Truncate if needed
            if input_ids.size(1) >= self.config["max_position_embeddings"]:
                input_ids = input_ids[:, -self.config["max_position_embeddings"]+1:]
            
            outputs = self(input_ids)
            logits = outputs["logits"][:, -1, :]
            
            # Apply repetition penalty
            if repetition_penalty != 1.0:
                for token_id in set(input_ids[0].tolist()):
                    logits[0, token_id] /= repetition_penalty
            
            # Apply temperature
            logits = logits / temperature
            
            # Top-k filtering
            if top_k > 0:
                indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
                logits[indices_to_remove] = float('-inf')
            
            # Top-p (nucleus) filtering
            sorted_logits, sorted_indices = torch.sort(logits, descending=True)
            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
            sorted_indices_to_remove = cumulative_probs > top_p
            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
            sorted_indices_to_remove[..., 0] = 0
            
            indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
            logits[indices_to_remove] = float('-inf')
            
            # Sample
            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            input_ids = torch.cat([input_ids, next_token], dim=1)
            
            # Check for EOS
            if eos_token_id is not None and next_token.item() == eos_token_id:
                break
        
        return input_ids