# rope_torch.py
import torch
import torch.nn as nn
import math

BATCH_SIZE = 16
SEQ_LEN = 256
HEAD_DIM = 128 
TOTAL_DIM = SEQ_LEN * HEAD_DIM

ROPE_THETA = 10000.0

def precompute_freqs(dim: int, theta: float = ROPE_THETA):
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
    t = torch.arange(SEQ_LEN)
    freqs = torch.outer(t, freqs).float()
    
    cos = torch.cos(freqs).to(torch.float32)
    sin = torch.sin(freqs).to(torch.float32)
    return cos.cuda(), sin.cuda()


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


class Model(nn.Module):

    def __init__(self, cos_cached, sin_cached):
        super().__init__()
        self.register_buffer('cos', cos_cached.unsqueeze(0))
        self.register_buffer('sin', sin_cached.unsqueeze(0))
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B_times_L, D = x.shape
        B = BATCH_SIZE
        L = SEQ_LEN

        x_reshaped = x.view(B, L, D) 

        x_pass, x_rot = x_reshaped.chunk(2, dim=-1)
        
        x_rot_half = rotate_half(x_rot)

        x_rot_comp = x_rot * self.cos
        x_rot_half_comp = x_rot_half * self.sin
        
        x_rot_result = x_rot_comp + x_rot_half_comp
        
        output = torch.cat((x_pass, x_rot_result), dim=-1)
        return output.view(B_times_L, D)


def get_inputs():
    x = torch.randn(BATCH_SIZE * SEQ_LEN, HEAD_DIM, dtype=torch.float32)
    return [x]

def get_init_inputs():
    cos, sin = precompute_freqs(HEAD_DIM, ROPE_THETA)
    return [cos, sin]