import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Apply rotary position embedding to incoming keys and update the KV cache.

    This isolates a common decode-time kernel: rotate the fresh key vectors for
    the current positions and store both keys and values into the running cache.
    """

    def __init__(self):
        super().__init__()

    def forward(
        self,
        k_new: torch.Tensor,
        v_new: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        cache_k: torch.Tensor,
        cache_v: torch.Tensor,
        positions: torch.Tensor,
    ) -> torch.Tensor:
        updated_k = cache_k.clone()
        updated_v = cache_v.clone()
        rot_dim = k_new.shape[-1]

        for b in range(k_new.shape[0]):
            for t in range(k_new.shape[1]):
                pos = int(positions[b, t].item())
                c = cos[pos]
                s = sin[pos]
                x = k_new[b, t]
                x_even = x[..., 0::2]
                x_odd = x[..., 1::2]
                rot_even = x_even * c - x_odd * s
                rot_odd = x_even * s + x_odd * c
                rotated = torch.empty_like(x)
                rotated[..., 0::2] = rot_even
                rotated[..., 1::2] = rot_odd
                updated_k[b, pos] = rotated
                updated_v[b, pos] = v_new[b, t]

        return torch.stack((updated_k, updated_v), dim=0)


batch_size = 2
cache_len = 128
update_len = 16
num_heads = 4
head_dim = 32


def _rope_tables(cache_len: int, half_dim: int):
    positions = torch.arange(cache_len, dtype=torch.float32).unsqueeze(1)
    freqs = torch.pow(10000.0, -torch.arange(half_dim, dtype=torch.float32) / half_dim).unsqueeze(0)
    angles = positions * freqs
    return torch.cos(angles), torch.sin(angles)


def get_inputs():
    p = popcorn_pri
    k_new = torch.randn(p.jitter_int(batch_size), p.jitter_int(update_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    v_new = torch.randn(p.jitter_int(batch_size), p.jitter_int(update_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    cos, sin = _rope_tables(p.jitter_int(cache_len), p.jitter_int(head_dim, align=8) // 2)
    cache_k = torch.randn(p.jitter_int(batch_size), p.jitter_int(cache_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    cache_v = torch.randn(p.jitter_int(batch_size), p.jitter_int(cache_len), p.jitter_int(num_heads), p.jitter_int(head_dim, align=8), dtype=torch.float32)
    base = torch.arange(p.jitter_int(update_len), dtype=torch.int32).unsqueeze(0).repeat(p.jitter_int(batch_size), 1)
    positions = (base + torch.tensor([[7], [31]], dtype=torch.int32)) % p.jitter_int(cache_len)
    return [k_new, v_new, cos, sin, cache_k, cache_v, positions]


def get_init_inputs():
    return []

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

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

__global__ void rope_update_kernel(
    const float* __restrict__ k_new,
    const float* __restrict__ v_new,
    const float* __restrict__ cos,
    const float* __restrict__ sin,
    const int32_t* __restrict__ positions,
    float* __restrict__ out_k,
    float* __restrict__ out_v,
    int batch_size,
    int update_len,
    int num_heads,
    int head_dim,
    int cache_len
) {
    int batch = blockIdx.x;
    int t = threadIdx.x;
    
    if (batch >= batch_size || t >= update_len) return;
    
    int pos = positions[batch * update_len + t];
    int half_dim = head_dim / 2;
    
    // Calculate strides
    long k_new_stride_b = update_len * num_heads * head_dim;
    long k_new_stride_t = num_heads * head_dim;
    long k_new_stride_h = head_dim;
    
    long cache_stride_b = cache_len * num_heads * head_dim;
    long cache_stride_p = num_heads * head_dim;
    long cache_stride_h = head_dim;
    
    long base_k_new = batch * k_new_stride_b + t * k_new_stride_t;
    long base_cache = batch * cache_stride_b + pos * cache_stride_p;
    
    // Copy and rotate k_new[batch, t] to out_k[batch, pos]
    for (int h = 0; h < num_heads; h++) {
        for (int d = 0; d < half_dim; d++) {
            float c = cos[pos * half_dim + d];
            float s = sin[pos * half_dim + d];
            
            float x_even = k_new[base_k_new + h * k_new_stride_h + d * 2];
            float x_odd = k_new[base_k_new + h * k_new_stride_h + d * 2 + 1];
            
            float rot_even = x_even * c - x_odd * s;
            float rot_odd = x_even * s + x_odd * c;
            
            out_k[base_cache + h * cache_stride_h + d * 2] = rot_even;
            out_k[base_cache + h * cache_stride_h + d * 2 + 1] = rot_odd;
        }
        
        // Copy v_new - copy ALL dimensions
        for (int d = 0; d < head_dim; d++) {
            out_v[base_cache + h * cache_stride_h + d] = 
                v_new[base_k_new + h * k_new_stride_h + d];
        }
    }
}

torch::Tensor forward_cuda(
    torch::Tensor k_new,
    torch::Tensor v_new,
    torch::Tensor cos,
    torch::Tensor sin,
    torch::Tensor cache_k,
    torch::Tensor cache_v,
    torch::Tensor positions
) {
    int batch_size = cache_k.size(0);
    int cache_len = cache_k.size(1);
    int num_heads = cache_k.size(2);
    int head_dim = cache_k.size(3);
    int update_len = k_new.size(1);
    
    // Clone caches to output
    auto out_k = cache_k.clone();
    auto out_v = cache_v.clone();
    
    // Ensure all tensors are contiguous
    auto k_new_c = k_new.contiguous();
    auto v_new_c = v_new.contiguous();
    auto cos_c = cos.contiguous();
    auto sin_c = sin.contiguous();
    auto positions_int = positions.to(torch::kInt32).contiguous();
    
    int threads_per_block = update_len;
    int num_blocks = batch_size;
    
    rope_update_kernel<<<num_blocks, threads_per_block>>>(
        k_new_c.data_ptr<float>(),
        v_new_c.data_ptr<float>(),
        cos_c.data_ptr<float>(),
        sin_c.data_ptr<float>(),
        positions_int.data_ptr<int32_t>(),
        out_k.data_ptr<float>(),
        out_v.data_ptr<float>(),
        batch_size,
        update_len,
        num_heads,
        head_dim,
        cache_len
    );
    
    return torch::stack({out_k, out_v}, 0);
}
"""

cpp_source = """
torch::Tensor forward_cuda(
    torch::Tensor k_new,
    torch::Tensor v_new,
    torch::Tensor cos,
    torch::Tensor sin,
    torch::Tensor cache_k,
    torch::Tensor cache_v,
    torch::Tensor positions
);
"""

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

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.op = rope_extension.forward_cuda
    
    def forward(self, k_new, v_new, cos, sin, cache_k, cache_v, positions):
        return self.op(k_new, v_new, cos, sin, cache_k, cache_v, positions)