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

cpp_src = r'''
#include <torch/extension.h>
torch::Tensor rope_update(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);
'''

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

template <typename PosT>
__global__ void rope_cache_kernel(
    const float* __restrict__ k_new,
    const float* __restrict__ v_new,
    const float* __restrict__ cos,
    const float* __restrict__ sin,
    const float* __restrict__ cache_k,
    const float* __restrict__ cache_v,
    const PosT* __restrict__ positions,
    float* __restrict__ out,
    int B, int U, int C, int H, int D)
{
    int64_t total_one = (int64_t)B * C * H * D;
    int64_t total = total_one * 2;
    int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * gridDim.x;
    int half = D >> 1;
    for (; idx < total; idx += stride) {
        bool is_v = idx >= total_one;
        int64_t j = is_v ? idx - total_one : idx;

        int d = j & (D - 1);
        int64_t q = j >> 5;
        int h = q & (H - 1);
        q >>= 2;
        int pos = q & (C - 1);
        int b = q >> 7;

        int first = (int)positions[b * U];
        int found_t = pos - first;
        if (found_t < 0) found_t += C;
        if (found_t >= U) found_t = -1;

        float val;
        if (found_t >= 0) {
            int64_t new_base = (((int64_t)b * U + found_t) * H + h) * D;
            if (is_v) {
                val = v_new[new_base + d];
            } else {
                int pair = d >> 1;
                float xe = k_new[new_base + (pair << 1)];
                float xo = k_new[new_base + (pair << 1) + 1];
                float c = cos[(int64_t)pos * half + pair];
                float s = sin[(int64_t)pos * half + pair];
                val = ((d & 1) == 0) ? (xe * c - xo * s) : (xe * s + xo * c);
            }
        } else {
            val = is_v ? cache_v[j] : cache_k[j];
        }
        out[idx] = val;
    }
}

torch::Tensor rope_update(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 B = (int)k_new.size(0);
    int U = (int)k_new.size(1);
    int H = (int)k_new.size(2);
    int D = (int)k_new.size(3);
    int C = (int)cache_k.size(1);
    auto out = torch::empty({2, B, C, H, D}, cache_k.options());
    int64_t total = (int64_t)2 * B * C * H * D;
    int threads = 256;
    int blocks = (int)((total + threads - 1) / threads);
    if (blocks > 256) blocks = 256;
    auto stream = at::cuda::getCurrentCUDAStream();
    if (positions.scalar_type() == at::kInt) {
        rope_cache_kernel<int><<<blocks, threads, 0, stream>>>(
            k_new.data_ptr<float>(), v_new.data_ptr<float>(), cos.data_ptr<float>(), sin.data_ptr<float>(),
            cache_k.data_ptr<float>(), cache_v.data_ptr<float>(), positions.data_ptr<int>(), out.data_ptr<float>(),
            B, U, C, H, D);
    } else if (positions.scalar_type() == at::kLong) {
        rope_cache_kernel<int64_t><<<blocks, threads, 0, stream>>>(
            k_new.data_ptr<float>(), v_new.data_ptr<float>(), cos.data_ptr<float>(), sin.data_ptr<float>(),
            cache_k.data_ptr<float>(), cache_v.data_ptr<float>(), positions.data_ptr<int64_t>(), out.data_ptr<float>(),
            B, U, C, H, D);
    } else {
        rope_cache_kernel<float><<<blocks, threads, 0, stream>>>(
            k_new.data_ptr<float>(), v_new.data_ptr<float>(), cos.data_ptr<float>(), sin.data_ptr<float>(),
            cache_k.data_ptr<float>(), cache_v.data_ptr<float>(), positions.data_ptr<float>(), out.data_ptr<float>(),
            B, U, C, H, D);
    }
    return out;
}
'''

_ext = load_inline(name='rope_cache_ext_final', cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=['rope_update'], extra_cuda_cflags=['-O3'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, k_new, v_new, cos, sin, cache_k, cache_v, positions):
        return _ext.rope_update(k_new, v_new, cos, sin, cache_k, cache_v, positions)