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
from torch.utils.cpp_extension import load_inline
import math

cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

__global__ void rope_kernel(
    const float* __restrict__ q,
    const float* __restrict__ k,
    float* __restrict__ q_out,
    float* __restrict__ k_out,
    const float* __restrict__ inv_freq,
    int batch_size, int num_heads, int seq_len, int dim
) {
    int half_dim = dim / 2;
    int total = batch_size * num_heads * seq_len * dim;
    
    for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; idx += blockDim.x * gridDim.x) {
        int d = idx % dim;
        int n = (idx / dim) % seq_len;
        int h = (idx / (dim * seq_len)) % num_heads;
        int b = idx / (dim * seq_len * num_heads);
        
        int base_idx = ((b * num_heads + h) * seq_len + n) * dim;
        
        int freq_idx = (d < half_dim) ? d : (d - half_dim);
        
        float inv_f = inv_freq[freq_idx];
        float angle = (float)n * inv_f;
        float cos_val, sin_val;
        sincosf(angle, &sin_val, &cos_val);
        
        if (d < half_dim) {
            float q_x1 = q[base_idx + d];
            float k_x1 = k[base_idx + d];
            float q_x2 = q[base_idx + d + half_dim];
            float k_x2 = k[base_idx + d + half_dim];
            q_out[base_idx + d] = q_x1 * cos_val - q_x2 * sin_val;
            k_out[base_idx + d] = k_x1 * cos_val - k_x2 * sin_val;
        } else {
            float q_x2 = q[base_idx + d];
            float k_x2 = k[base_idx + d];
            float q_x1 = q[base_idx + d - half_dim];
            float k_x1 = k[base_idx + d - half_dim];
            q_out[base_idx + d] = q_x2 * cos_val + q_x1 * sin_val;
            k_out[base_idx + d] = k_x2 * cos_val + k_x1 * sin_val;
        }
    }
}

torch::Tensor apply_rope_cuda(
    torch::Tensor q,
    torch::Tensor k,
    torch::Tensor inv_freq,
    int64_t max_seq_len,
    double base
) {
    auto q_out = torch::empty_like(q);
    auto k_out = torch::empty_like(k);
    
    int batch_size = q.size(0);
    int num_heads = q.size(1);
    int seq_len = q.size(2);
    int dim = q.size(3);
    
    int total = batch_size * num_heads * seq_len * dim;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    blocks = std::min(blocks, 65535);
    
    rope_kernel<<<blocks, threads>>>(
        q.data_ptr<float>(),
        k.data_ptr<float>(),
        q_out.data_ptr<float>(),
        k_out.data_ptr<float>(),
        inv_freq.data_ptr<float>(),
        batch_size, num_heads, seq_len, dim
    );
    
    return torch::stack({q_out, k_out}, 0);
}
"""

cpp_source = """
torch::Tensor apply_rope_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor inv_freq, int64_t max_seq_len, double base);
"""

rope_extension = load_inline(
    name='rope_extension_v5',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['apply_rope_cuda'],
    extra_cuda_cflags=['-O3'],
    verbose=False
)


class ModelNew(nn.Module):
    """
    Rotary Position Embedding (RoPE) applied to protein residue sequences.
    Optimized CUDA implementation.
    """

    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
        self.base = base

    def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple:
        result = rope_extension.apply_rope_cuda(
            q.contiguous(), 
            k.contiguous(), 
            self.inv_freq,
            self.max_seq_len,
            self.base
        )
        return result[0], result[1]