import torch
import torch.nn as nn

class RoPE(nn.Module):
    
    def __init__(self, d_head, context_len, dtype = torch.float16):
       
        super().__init__()
        
        self.context_len = context_len
        
        position = torch.arange(start = 0, end = context_len, dtype = dtype)
        div_term = 10000 ** (torch.arange(start = 0, end = d_head, step = 2, dtype = torch.float16) / d_head)
        div_term = torch.repeat_interleave(div_term, repeats=2, dim=-1)
        rope_cos = torch.cos(position / div_term) 
        rope_sin = torch.sin(position / div_term)
        
        self.register_buffer("rope_cos", rope_cos)
        self.register_buffer("rope_sin", rope_sin)
        
    def forward(self, x): # x is batch_size, n_heads, seq_len, d_head
        
        cos = self.rope_cos[:x.shape[2]]
        sin = self.rope_sin[:x.shape[2]]
        
        cos = cos.unsqueeze(0).unsqueeze(0)
        sin = sin.unsqueeze(0).unsqueeze(0)
        
        x_even = x[..., 0::2]
        x_odd = x[..., 1::2]
        cos_even = cos[..., 0::2]
        sin_even = sin[..., 0::2]  
        
        rotated_x_even = x_even * cos_even - x_odd * sin_even
        rotated_x_odd = x_even * sin_even + x_odd * cos_even
        
        rotated = torch.stack([rotated_x_even, rotated_x_odd], dim = -1)
        rotated = rotated.flatten(-2) 
        return rotated