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

_cpp_src = r"""
#include <torch/extension.h>
#include <tuple>
std::tuple<torch::Tensor, torch::Tensor> rope_forward(torch::Tensor q, torch::Tensor k, torch::Tensor cos, torch::Tensor sin);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <tuple>

__global__ void rope64_kernel(const float* __restrict__ q,
                              const float* __restrict__ k,
                              const float* __restrict__ cosv,
                              const float* __restrict__ sinv,
                              float* __restrict__ qo,
                              float* __restrict__ ko,
                              int BH, int N) {
    int i = threadIdx.x;
    int bh = blockIdx.y * 8 + threadIdx.y;
    int n = blockIdx.x;
    if (bh < BH) {
        int base = (bh * N + n) * 64 + i;
        float c = cosv[n * 64 + i];
        float s = sinv[n * 64 + i];
        float q1 = q[base], q2 = q[base + 32];
        float k1 = k[base], k2 = k[base + 32];
        qo[base]      = q1 * c - q2 * s;
        qo[base + 32] = q2 * c + q1 * s;
        ko[base]      = k1 * c - k2 * s;
        ko[base + 32] = k2 * c + k1 * s;
    }
}

__global__ void rope_pair_kernel(const float* __restrict__ q,
                                 const float* __restrict__ k,
                                 const float* __restrict__ cosv,
                                 const float* __restrict__ sinv,
                                 float* __restrict__ qo,
                                 float* __restrict__ ko,
                                 int total, int N, int D, int half) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    for (; idx < total; idx += stride) {
        int i = idx % half;
        int t = idx / half;
        int n = t % N;
        int bh = t / N;
        int base = (bh * N + n) * D + i;
        float c = cosv[n * D + i];
        float s = sinv[n * D + i];
        float q1 = q[base], q2 = q[base + half];
        float k1 = k[base], k2 = k[base + half];
        qo[base]        = q1 * c - q2 * s;
        qo[base + half] = q2 * c + q1 * s;
        ko[base]        = k1 * c - k2 * s;
        ko[base + half] = k2 * c + k1 * s;
    }
}

std::tuple<torch::Tensor, torch::Tensor> rope_forward(torch::Tensor q, torch::Tensor k, torch::Tensor cos, torch::Tensor sin) {
    int64_t B64 = q.size(0), H64 = q.size(1), N64 = q.size(2), D64 = q.size(3);
    auto qo = torch::empty({B64, H64, N64, D64}, q.options());
    auto ko = torch::empty({B64, H64, N64, D64}, q.options());
    int B = (int)B64, H = (int)H64, N = (int)N64, D = (int)D64;
    int BH = B * H;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (D == 64) {
        dim3 block(32, 8, 1);
        dim3 grid(N, (BH + 7) / 8, 1);
        rope64_kernel<<<grid, block, 0, stream>>>(q.data_ptr<float>(), k.data_ptr<float>(), cos.data_ptr<float>(), sin.data_ptr<float>(), qo.data_ptr<float>(), ko.data_ptr<float>(), BH, N);
    } else {
        int half = D >> 1;
        int total = BH * N * half;
        int threads = 256;
        int blocks = (total + threads - 1) / threads;
        if (blocks > 4096) blocks = 4096;
        rope_pair_kernel<<<blocks, threads, 0, stream>>>(q.data_ptr<float>(), k.data_ptr<float>(), cos.data_ptr<float>(), sin.data_ptr<float>(), qo.data_ptr<float>(), ko.data_ptr<float>(), total, N, D, half);
    }
    return std::make_tuple(qo, ko);
}
"""

_rope_ext = load_inline(name="rope_ext_pair_final", cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=["rope_forward"], extra_cflags=["-O3"], extra_cuda_cflags=["-O3"], verbose=False)

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))
        t = torch.arange(max_seq_len, dtype=inv_freq.dtype)
        freqs = torch.outer(t, inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        self.register_buffer("cos_cached", emb.cos().contiguous())
        self.register_buffer("sin_cached", emb.sin().contiguous())
        self.max_seq_len = max_seq_len
        self.dim = dim

    def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple:
        return _rope_ext.rope_forward(q, k, self.cos_cached, self.sin_cached)