import torch
import torch.nn as nn
import math


class Model(nn.Module):
    """
    Rotary Position Embedding (RoPE) applied to protein residue sequences.
    Used in ESM-2 and similar protein language models.  Encodes absolute
    position via rotation of query/key pairs in 2-D subspaces, preserving
    relative-position information in dot-product attention.
    """

    def __init__(self, dim, max_seq_len=2048, base=10000.0):
        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
        self.dim = dim

    def _build_cos_sin(self, seq_len: int, device: torch.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()

    @staticmethod
    def _rotate_half(x: torch.Tensor) -> torch.Tensor:
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat([-x2, x1], dim=-1)

    def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple:
        """
        Args:
            q: (B, H, N, D)  – queries
            k: (B, H, N, D)  – keys
        Returns:
            (q_rot, k_rot) each (B, H, N, D) with RoPE applied
        """
        N = q.shape[2]
        cos, sin = self._build_cos_sin(N, q.device)  # (N, D)
        cos = cos.unsqueeze(0).unsqueeze(0)
        sin = sin.unsqueeze(0).unsqueeze(0)
        q_rot = q * cos + self._rotate_half(q) * sin
        k_rot = k * cos + self._rotate_half(k) * sin
        return q_rot, k_rot


dim = 64
num_heads = 8
seq_len = 256
batch_size = 4


def get_inputs():
    p = popcorn_pri
    bs = p.jitter_int(batch_size)
    nh = p.jitter_int(num_heads)
    sl = p.jitter_int(seq_len)
    q = torch.randn(bs, nh, sl, dim)
    k = torch.randn(bs, nh, sl, dim)
    return [q, k]


def get_init_inputs():
    return [dim]

import torch
import torch.nn as nn
import torch.utils.cpp_extension as cpp_extension

# CUDA kernel source code
cuda_source = '''
#include <cuda_runtime.h>
#include <cuda_fp16.h>

__global__ void rope_cuda_kernel(
    const float* q,
    const float* k,
    float* q_rot,
    float* k_rot,
    const float* cos,
    const float* sin,
    int batch_size,
    int num_heads,
    int seq_len,
    int dim
) {
    int b = blockIdx.x;
    int h = blockIdx.y;
    int n = threadIdx.x + blockIdx.z * blockDim.x;
    if (n >= seq_len) return;

    int offset_q = b * num_heads * seq_len * dim + h * seq_len * dim + n * dim;
    int offset_cos = n * dim;

    for (int d = 0; d < dim; d += 2) {
        float q0 = q[offset_q + d];
        float q1 = q[offset_q + d + 1];
        float cos_val = cos[offset_cos + d];
        float sin_val = sin[offset_cos + d];
        q_rot[offset_q + d] = q0 * cos_val - q1 * sin_val;
        q_rot[offset_q + d + 1] = q0 * sin_val + q1 * cos_val;

        float k0 = k[offset_q + d];
        float k1 = k[offset_q + d + 1];
        k_rot[offset_q + d] = k0 * cos_val - k1 * sin_val;
        k_rot[offset_q + d + 1] = k0 * sin_val + k1 * cos_val;
    }
}

__global__ void rope_cuda_kernel_half(
    const half* q,
    const half* k,
    half* q_rot,
    half* k_rot,
    const half* cos,
    const half* sin,
    int batch_size,
    int num_heads,
    int seq_len,
    int dim
) {
    int b = blockIdx.x;
    int h = blockIdx.y;
    int n = threadIdx.x + blockIdx.z * blockDim.x;
    if (n >= seq_len) return;

    int offset_q = b * num_heads * seq_len * dim + h * seq_len * dim + n * dim;
    int offset_cos = n * dim;

    for (int d = 0; d < dim; d += 2) {
        half q0 = q[offset_q + d];
        half q1 = q[offset_q + d + 1];
        half cos_val = cos[offset_cos + d];
        half sin_val = sin[offset_cos + d];
        q_rot[offset_q + d] = __hmul(q0, cos_val) - __hmul(q1, sin_val);
        q_rot[offset_q + d + 1] = __hmul(q0, sin_val) + __hmul(q1, cos_val);

        half k0 = k[offset_q + d];
        half k1 = k[offset_q + d + 1];
        k_rot[offset_q + d] = __hmul(k0, cos_val) - __hmul(k1, sin_val);
        k_rot[offset_q + d + 1] = __hmul(k0, sin_val) + __hmul(k1, cos_val);
    }
}

void rope_cuda(
    const void* q,
    const void* k,
    void* q_rot,
    void* k_rot,
    const void* cos,
    const void* sin,
    int batch_size,
    int num_heads,
    int seq_len,
    int dim,
    bool use_half
) {
    dim3 grid(batch_size, num_heads, (seq_len + 255) / 256);
    dim3 block(256);
    if (use_half) {
        rope_cuda_kernel_half<<<grid, block>>>((const half*)q, (const half*)k, (half*)q_rot, (half*)k_rot, (const half*)cos, (const half*)sin, batch_size, num_heads, seq_len, dim);
    } else {
        rope_cuda_kernel<<<grid, block>>>((const float*)q, (const float*)k, (float*)q_rot, (float*)k_rot, (const float*)cos, (const float*)sin, batch_size, num_heads, seq_len, dim);
    }
}
'''

# Load CUDA kernel
rope_cuda_module = cpp_extension.load_inline(
    name='rope_cuda',
    cuda_sources=cuda_source,
    functions=['rope_cuda'],
    extra_cuda_cflags=['-O3']
)

class ModelNew(nn.Module):
    def __init__(self, dim, max_seq_len=2048, base=10000.0):
        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
        self.dim = dim

    def _build_cos_sin(self, seq_len: int, device: torch.device, dtype: torch.dtype):
        t = torch.arange(seq_len, device=device, dtype=dtype)
        freqs = torch.outer(t, self.inv_freq.to(dtype))
        emb = torch.cat([freqs, freqs], dim=-1)
        return emb.cos(), emb.sin()

    def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple:
        N = q.shape[2]
        use_half = q.dtype == torch.float16
        cos, sin = self._build_cos_sin(N, q.device, q.dtype)
        cos = cos.unsqueeze(0).unsqueeze(0)
        sin = sin.unsqueeze(0).unsqueeze(0)
        q_rot = torch.empty_like(q)
        k_rot = torch.empty_like(k)
        rope_cuda_module.rope_cuda(
            q.data_ptr(),
            k.data_ptr(),
            q_rot.data_ptr(),
            k_rot.data_ptr(),
            cos.data_ptr(),
            sin.data_ptr(),
            *q.shape[:3],
            q.shape[-1],
            use_half
        )
        return q_rot, k_rot