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 <vector>

std::vector<torch::Tensor> rope_forward(torch::Tensor q, torch::Tensor k, double base);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#include <math.h>

__global__ void rope_f32_kernel(const float* __restrict__ q,
                                const float* __restrict__ k,
                                float* __restrict__ qo,
                                float* __restrict__ ko,
                                long total_pairs,
                                int Dhalf,
                                int N,
                                int D,
                                float log_base_over_D) {
    long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total_pairs) return;

    int j = (int)(idx % Dhalf);
    long t = idx / Dhalf;
    int n = (int)(t % N);
    long base_idx = (t / N) * (long)N * D + (long)n * D + j;

    // inv_freq[j] = exp(-(2*j/dim)*log(base));  angle = n * inv_freq[j]
    float invf = __expf(-((float)j) * log_base_over_D);
    float angle = (float)n * invf;
    float s, c;
    __sincosf(angle, &s, &c);

    long i1 = base_idx;
    long i2 = base_idx + Dhalf;

    float q1 = q[i1];
    float q2 = q[i2];
    float k1 = k[i1];
    float k2 = k[i2];

    qo[i1] = q1 * c - q2 * s;
    qo[i2] = q2 * c + q1 * s;
    ko[i1] = k1 * c - k2 * s;
    ko[i2] = k2 * c + k1 * s;
}

std::vector<torch::Tensor> rope_forward(torch::Tensor q, torch::Tensor k, double base) {
    auto qo = torch::empty_like(q);
    auto ko = torch::empty_like(k);

    const int D = (int)q.size(3);
    const int Dhalf = D >> 1;
    const int N = (int)q.size(2);
    const long B = q.size(0);
    const long H = q.size(1);
    const long total_pairs = B * H * (long)N * Dhalf;

    if (total_pairs > 0) {
        const int threads = 256;
        const int blocks = (int)((total_pairs + threads - 1) / threads);
        float log_base_over_D = (float)(2.0 * log(base) / (double)D);
        rope_f32_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
            q.data_ptr<float>(), k.data_ptr<float>(), qo.data_ptr<float>(), ko.data_ptr<float>(),
            total_pairs, Dhalf, N, D, log_base_over_D);
    }
    return {qo, ko};
}
'''

_ext = load_inline(
    name='rope_inline_ext_fast_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['rope_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_cflags=['-O3'],
    with_cuda=True,
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, dim, max_seq_len=2048, base=10000.0):
        super().__init__()
        self.dim = dim
        self.max_seq_len = max_seq_len
        self.base = float(base)
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)

    def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple:
        out = _ext.rope_forward(q, k, self.base)
        return out[0], out[1]