import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import asdict
from .config import TinyConfig

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__(); self.eps=eps; self.weight=nn.Parameter(torch.ones(dim))
    def forward(self, x):
        return self.weight * x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

def rotate_half(x):
    x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
    return torch.cat((-x2, x1), dim=-1)

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000.0):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        t = torch.arange(max_position_embeddings, dtype=torch.float)
        freqs = torch.einsum('i,j->ij', t, inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer('cos_cached', emb.cos()[None, None, :, :], persistent=False)
        self.register_buffer('sin_cached', emb.sin()[None, None, :, :], persistent=False)
    def forward(self, q, k, seq_len):
        cos = self.cos_cached[:, :, :seq_len, :].to(q.device, q.dtype)
        sin = self.sin_cached[:, :, :seq_len, :].to(q.device, q.dtype)
        return (q*cos + rotate_half(q)*sin), (k*cos + rotate_half(k)*sin)

class CausalSelfAttention(nn.Module):
    def __init__(self, cfg: TinyConfig):
        super().__init__()
        assert cfg.hidden_size % cfg.num_attention_heads == 0
        assert cfg.num_attention_heads % cfg.num_key_value_heads == 0
        self.nh=cfg.num_attention_heads; self.nkv=cfg.num_key_value_heads
        self.hd=cfg.hidden_size//cfg.num_attention_heads
        self.q_proj=nn.Linear(cfg.hidden_size, self.nh*self.hd, bias=cfg.attention_bias)
        self.k_proj=nn.Linear(cfg.hidden_size, self.nkv*self.hd, bias=cfg.attention_bias)
        self.v_proj=nn.Linear(cfg.hidden_size, self.nkv*self.hd, bias=cfg.attention_bias)
        self.o_proj=nn.Linear(cfg.hidden_size, cfg.hidden_size, bias=cfg.attention_bias)
        self.rotary=RotaryEmbedding(self.hd, cfg.max_position_embeddings, cfg.rope_theta)
        self.dropout=cfg.dropout
    def forward(self, x):
        B,T,C=x.shape
        q=self.q_proj(x).view(B,T,self.nh,self.hd).transpose(1,2)
        k=self.k_proj(x).view(B,T,self.nkv,self.hd).transpose(1,2)
        v=self.v_proj(x).view(B,T,self.nkv,self.hd).transpose(1,2)
        q,k=self.rotary(q,k,T)
        if self.nkv != self.nh:
            repeat = self.nh // self.nkv
            k = k.repeat_interleave(repeat, dim=1)
            v = v.repeat_interleave(repeat, dim=1)
        y=F.scaled_dot_product_attention(q,k,v,dropout_p=self.dropout if self.training else 0.0,is_causal=True)
        y=y.transpose(1,2).contiguous().view(B,T,C)
        return self.o_proj(y)

class SwiGLU(nn.Module):
    def __init__(self, cfg: TinyConfig):
        super().__init__()
        self.gate_proj=nn.Linear(cfg.hidden_size,cfg.intermediate_size,bias=cfg.mlp_bias)
        self.up_proj=nn.Linear(cfg.hidden_size,cfg.intermediate_size,bias=cfg.mlp_bias)
        self.down_proj=nn.Linear(cfg.intermediate_size,cfg.hidden_size,bias=cfg.mlp_bias)
    def forward(self,x):
        return self.down_proj(F.silu(self.gate_proj(x))*self.up_proj(x))

class Block(nn.Module):
    def __init__(self,cfg):
        super().__init__(); self.input_norm=RMSNorm(cfg.hidden_size,cfg.rms_norm_eps); self.attn=CausalSelfAttention(cfg); self.post_norm=RMSNorm(cfg.hidden_size,cfg.rms_norm_eps); self.mlp=SwiGLU(cfg)
    def forward(self,x):
        x=x+self.attn(self.input_norm(x)); x=x+self.mlp(self.post_norm(x)); return x

class TinyLlamaForCausalLM(nn.Module):
    def __init__(self,cfg:TinyConfig):
        super().__init__(); self.config=cfg
        self.embed_tokens=nn.Embedding(cfg.vocab_size,cfg.hidden_size)
        self.layers=nn.ModuleList([Block(cfg) for _ in range(cfg.num_hidden_layers)])
        self.norm=RMSNorm(cfg.hidden_size,cfg.rms_norm_eps)
        self.lm_head=nn.Linear(cfg.hidden_size,cfg.vocab_size,bias=False)
        if cfg.tie_word_embeddings:
            self.lm_head.weight=self.embed_tokens.weight
        self.apply(self._init_weights)
    def _init_weights(self,m):
        if isinstance(m, nn.Linear): nn.init.normal_(m.weight, mean=0.0, std=0.02)
        if isinstance(m, nn.Embedding): nn.init.normal_(m.weight, mean=0.0, std=0.02)
    def forward(self,input_ids,labels=None,loss_mask=None):
        x=self.embed_tokens(input_ids)
        for layer in self.layers: x=layer(x)
        logits=self.lm_head(self.norm(x))
        loss=None
        if labels is not None:
            shift_logits=logits[:,:-1,:].contiguous(); shift_labels=labels[:,1:].contiguous()
            per=F.cross_entropy(shift_logits.view(-1,shift_logits.size(-1)),shift_labels.view(-1),reduction='none')
            if loss_mask is not None:
                mask=loss_mask[:,1:].contiguous().view(-1).float(); loss=(per*mask).sum()/mask.sum().clamp_min(1.0)
            else: loss=per.mean()
        return {'loss':loss,'logits':logits}
    def num_parameters(self): return sum(p.numel() for p in self.parameters())
    def save_config(self,path):
        import json; open(path,'w').write(json.dumps(asdict(self.config),indent=2))