#! /usr/bin/env python
"""
SmollmV2 model implementation
Author: Shilpaj Bhalerao
Date: 2025-01-19
"""
# Third-Party Imports
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

# Local Imports
from config import SmollmConfig, RoPEConfig


class RoPEAttention:
    """
    Rotary Position Embedding attention with support for different Q/K dimensions
    """
    def __init__(self, head_dim, kv_dim, base=RoPEConfig.base):
        """
        Initialize rotary embeddings
        Args:
            head_dim: Dimension of query head
            kv_dim: Dimension of key/value head
            base: Base for the angle calculations (default: 10000)
        """
        super().__init__()
        
        # Generate theta parameter for rotary embeddings for both Q and K dimensions
        inv_freq_k = 1.0 / (base ** (torch.arange(0, kv_dim, 2).float() / kv_dim))
        self.register_buffer('inv_freq_k', inv_freq_k)
        
        self.head_dim = head_dim
        self.kv_dim = kv_dim
        self.seq_len_cached = None
        self.cos_cached = None
        self.sin_cached = None

    def _update_cos_sin_cache(self, x, seq_len):
        """Update cached cos and sin values for given sequence length"""
        if seq_len != self.seq_len_cached:
            self.seq_len_cached = seq_len
            t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq_k)
            freqs = torch.einsum('i,j->ij', t, self.inv_freq_k)
            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
            
            self.cos_cached = emb.cos()[None, None, :, :]
            self.sin_cached = emb.sin()[None, None, :, :]

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

    def __call__(self, q, k):
        """
        Apply rotary embeddings to input queries and keys
        Args:
            q: Query tensor of shape (batch, n_head, seq_len, head_dim)
            k: Key tensor of shape (batch, n_head, seq_len, kv_dim)
        Returns:
            q_rot: Rotated query tensor
            k_rot: Rotated key tensor
        """
        seq_len = q.shape[2]
        self._update_cos_sin_cache(k, seq_len)
        
        # Apply rotary embeddings to keys
        k_cos = self.cos_cached[..., :self.kv_dim]
        k_sin = self.sin_cached[..., :self.kv_dim]
        k_rot = (k * k_cos) + (self._rotate_half(k) * k_sin)
        
        # For queries, we only apply rotation to the part that interacts with keys
        q_part = q[..., :self.kv_dim]
        q_cos = self.cos_cached[..., :self.kv_dim]
        q_sin = self.sin_cached[..., :self.kv_dim]
        q_rot_part = (q_part * q_cos) + (self._rotate_half(q_part) * q_sin)
        
        # Combine rotated part with unrotated parts for query
        q_rot = torch.cat([q_rot_part, q[..., self.kv_dim:]], dim=-1)
        
        return q_rot, k_rot

    def register_buffer(self, name, tensor):
        """Helper function to register a buffer"""
        setattr(self, name, tensor)


class CausalSelfAttention(nn.Module):
    """
    Causal self-attention mechanism with reduced KV dimensions and RoPE
    """
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        
        # Calculate dimensions
        self.head_dim = config.n_embd // config.n_head  # 576/9 = 64
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        
        # Make kv_dim divisible by n_head (189 is closest to 192 that's divisible by 9)
        self.kv_dim = 189  # 189 = 9 * 21, closest to 192 that's divisible by 9
        self.kv_dim_per_head = self.kv_dim // self.n_head  # 21
        
        # Separate projections with reduced dimensions for k,v
        self.q_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
        self.k_proj = nn.Linear(config.n_embd, self.kv_dim, bias=False)  # 189 dimensions
        self.v_proj = nn.Linear(config.n_embd, self.kv_dim, bias=False)  # 189 dimensions
        
        # output projection
        self.o_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
        
        # rotary embeddings
        self.rope = RoPEAttention(self.head_dim, self.kv_dim_per_head)

    def forward(self, x):
        B, T, C = x.size()
        
        # calculate query, key, values
        q = self.q_proj(x)
        k = self.k_proj(x)
        v = self.v_proj(x)
        
        # reshape with exact dimensions
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.kv_dim_per_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.kv_dim_per_head).transpose(1, 2)

        # apply rotary embeddings
        q, k = self.rope(q, k)

        # pad k and v to match q dimension for attention
        k_pad = torch.zeros_like(q)
        v_pad = torch.zeros_like(q)
        k_pad[..., :self.kv_dim_per_head] = k
        v_pad[..., :self.kv_dim_per_head] = v

        # flash attention
        y = F.scaled_dot_product_attention(q, k_pad, v_pad, is_causal=True)

        # reshape back
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        
        # output projection
        y = self.o_proj(y)
        return y


class MLP(nn.Module):
    """
    MLP (Multi-Layer Perceptron) layer with gate/up/down projection structure
    """
    def __init__(self, config):
        super().__init__()
        hidden_dim = int(config.n_embd * config.mlp_ratio) - 1
        self.gate_proj = nn.Linear(config.n_embd, hidden_dim, bias=False)
        self.up_proj = nn.Linear(config.n_embd, hidden_dim, bias=False)
        self.down_proj = nn.Linear(hidden_dim, config.n_embd, bias=False)
        self.down_proj.NANOGPT_SCALE_INIT = 1

    def forward(self, x):
        # SwiGLU activation as used in PaLM, Llama, etc.
        gate = self.gate_proj(x)
        up = self.up_proj(x)
        x = F.silu(gate) * up
        x = self.down_proj(x)
        return x


class Block(nn.Module):
    """
    Transformer block
    """
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd, bias=False)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd, bias=False)
        self.mlp = MLP(config)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x


class SmollmV2(nn.Module):
    """
    SmollmV2 model
    """
    def __init__(self, config=SmollmConfig()):
        super().__init__()
        self.config = config

        self.transformer = nn.ModuleDict(dict(
            wte = nn.Embedding(config.vocab_size, config.n_embd),
            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = nn.LayerNorm(config.n_embd, bias=False),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        # weight sharing
        self.transformer.wte.weight = self.lm_head.weight

        # weight initialization
        self.apply(self._init_weights)
        
        # Compile the model if torch version supports it
        if hasattr(torch, 'compile'):
            self.forward = torch.compile(self.forward)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            std = 0.02
            if hasattr(module, 'NANGPT_SCALE_INIT'):
                std *= (2 * self.config.n_layer) ** -0.5
            torch.nn.init.normal_(module.weight, mean = 0.0, std = std)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std = 0.04)

    def forward(self, idx, targets=None):
        # idx is of shape (B, T)
        B, T = idx.size()
        assert T <= self.config.block_size, f"Cannot forward sequence of length {T}, block size is only {self.config.block_size}"
        # forward the token and posisition embeddings
        tok_emb = self.transformer.wte(idx) # token embeddings of shape (B, T, n_embd)
        x = tok_emb
        # forward the blocks of the transformer
        for block in self.transformer.h:
            x = block(x)
        # forward the final layernorm and the classifier
        x = self.transformer.ln_f(x)
        logits = self.lm_head(x) # (B, T, vocab_size)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss