import math

import torch
import torch.nn as nn
import torch.nn.functional as F

# Default ModelNew is PyTorch-only. run_and_check: use check_kernel=False with backend=cuda
# until the candidate file contains a real CUDA kernel (see kernel_static_checker).

class Model(nn.Module):
    """
    Single-layer multi-head cross-attention with RoPE on Q and K, attention
    dropout, and output projection—intended as a fusion target for robotics
    stacks (e.g. proprio / action queries attending to vision or memory).

    Pipeline (conceptually one fused kernel): linear Q from query stream,
    linear K/V from key–value stream → reshape heads → apply RoPE separately
    along query positions (T_q) and key positions (T_k) → scaled dot-product,
    softmax, dropout on weights → aggregate values → output projection.
    """

    def __init__(
        self,
        dim: int,
        num_heads: int,
        dropout: float = 0.0,
        rope_base: float = 10000.0,
    ):
        super().__init__()
        assert dim % num_heads == 0
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        assert self.head_dim % 2 == 0, "RoPE requires an even head dimension"

        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)
        self.out_proj = nn.Linear(dim, dim, bias=False)
        self.dropout = nn.Dropout(dropout)
        self.scale = 1.0 / math.sqrt(self.head_dim)

        inv_freq = 1.0 / (
            rope_base ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim)
        )
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    @staticmethod
    def _rotate_half(x: torch.Tensor) -> torch.Tensor:
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat([-x2, x1], dim=-1)

    def _rope(self, x: torch.Tensor, seq_len: int) -> torch.Tensor:
        """Apply RoPE along the sequence dimension (last-but-one after head). x: (B, H, L, d)."""
        t = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(t, self.inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        cos = emb.cos().view(1, 1, seq_len, self.head_dim)
        sin = emb.sin().view(1, 1, seq_len, self.head_dim)
        return x * cos + self._rotate_half(x) * sin

    def forward(self, x_query: torch.Tensor, x_kv: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x_query: (B, T_q, dim) — e.g. robot or decoder tokens
            x_kv:    (B, T_k, dim) — e.g. encoder / memory tokens
        Returns:
            (B, T_q, dim)
        """
        B, T_q, _ = x_query.shape
        _, T_k, _ = x_kv.shape
        h, d = self.num_heads, self.head_dim

        q = self.q_proj(x_query).view(B, T_q, h, d).transpose(1, 2)
        k = self.k_proj(x_kv).view(B, T_k, h, d).transpose(1, 2)
        v = self.v_proj(x_kv).view(B, T_k, h, d).transpose(1, 2)

        q = self._rope(q, T_q)
        k = self._rope(k, T_k)

        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
        attn = F.softmax(attn, dim=-1)
        attn = self.dropout(attn)
        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).contiguous().view(B, T_q, self.dim)
        return self.out_proj(out)


class ModelNew(nn.Module):
    """
    KernelBench candidate entry point (`eval_kernel_against_ref` / `run_and_check`).
    Replace this class with a fused CUDA/Triton (etc.) implementation; the default
    delegates to `Model` so you can smoke-test the harness with the same file for
    both `ref_arch_src_path` and `kernel_src_path`.
    """

    def __init__(
        self,
        dim: int,
        num_heads: int,
        dropout: float = 0.0,
        rope_base: float = 10000.0,
    ):
        super().__init__()
        self._impl = Model(dim, num_heads, dropout, rope_base)

    def forward(self, x_query: torch.Tensor, x_kv: torch.Tensor) -> torch.Tensor:
        return self._impl(x_query, x_kv)


batch_size = 4
seq_len_q = 64
seq_len_kv = 128
dim = 256
num_heads = 8
# Use 0.0 so `run_and_check` / multi-trial correctness matches: the harness does not
# re-seed RNG between the reference forward and `ModelNew` forward, so attention
# dropout > 0 would make outputs differ even for identical math.
dropout_p = 0.0


def get_inputs():
    p = popcorn_pri
    bs = p.jitter_int(batch_size)
    tq = p.jitter_int(seq_len_q)
    tk = p.jitter_int(seq_len_kv)
    x_query = torch.randn(bs, tq, dim)
    x_kv = torch.randn(bs, tk, dim)
    return [x_query, x_kv]


def get_init_inputs():
    return [dim, num_heads, dropout_p]

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Compile the CUDA kernel
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

// RoPE kernel: applies rotary position embeddings
// x: (B, H, L, d) - input tensor
// cos, sin: (L, d/2) - position embeddings
// out: (B, H, L, d) - output tensor
__global__ void rope_kernel(
    const float* __restrict__ x_ptr,
    const float* __restrict__ cos_ptr,
    const float* __restrict__ sin_ptr,
    float* __restrict__ out_ptr,
    int B, int H, int L, int d
) {
    int half_d = d / 2;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * H * L * half_d;
    
    if (idx < total) {
        int b = idx / (H * L * half_d);
        int h = (idx / (L * half_d)) % H;
        int l = (idx / half_d) % L;
        int i = idx % half_d;
        
        int base_idx = b * H * L * d + h * L * d + l * d;
        float x1 = x_ptr[base_idx + i];
        float x2 = x_ptr[base_idx + i + half_d];
        float c = cos_ptr[l * half_d + i];
        float s = sin_ptr[l * half_d + i];
        
        out_ptr[base_idx + i] = x1 * c - x2 * s;
        out_ptr[base_idx + i + half_d] = x1 * s + x2 * c;
    }
}

// Softmax kernel for attention weights
// attn: (B, H, T_q, T_k) - attention scores
// scale: scaling factor
__global__ void softmax_kernel(
    float* __restrict__ attn_ptr,
    int B, int H, int T_q, int T_k,
    float scale
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * H * T_q;
    
    if (idx < total) {
        int b = idx / (H * T_q);
        int h = (idx / T_q) % H;
        int q_idx = idx % T_q;
        
        int row_offset = b * H * T_q * T_k + h * T_q * T_k + q_idx * T_k;
        
        // Find max for numerical stability
        float max_val = -INFINITY;
        for (int k_idx = 0; k_idx < T_k; k_idx++) {
            float val = attn_ptr[row_offset + k_idx] * scale;
            if (val > max_val) max_val = val;
        }
        
        // Compute exp and sum
        float sum = 0.0f;
        for (int k_idx = 0; k_idx < T_k; k_idx++) {
            float val = expf(attn_ptr[row_offset + k_idx] * scale - max_val);
            attn_ptr[row_offset + k_idx] = val;
            sum += val;
        }
        
        // Normalize
        float inv_sum = 1.0f / sum;
        for (int k_idx = 0; k_idx < T_k; k_idx++) {
            attn_ptr[row_offset + k_idx] *= inv_sum;
        }
    }
}

// Attention matrix multiplication: Q @ K^T
__global__ void qk_kernel(
    const float* __restrict__ q_ptr,
    const float* __restrict__ k_ptr,
    float* __restrict__ attn_ptr,
    int B, int H, int T_q, int T_k, int d
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * H * T_q * T_k;
    
    if (idx < total) {
        int b = idx / (H * T_q * T_k);
        int h = (idx / (T_q * T_k)) % H;
        int q_idx = (idx / T_k) % T_q;
        int kv_idx = idx % T_k;
        
        int q_offset = b * H * T_q * d + h * T_q * d + q_idx * d;
        int k_offset = b * H * T_k * d + h * T_k * d + kv_idx * d;
        
        float sum = 0.0f;
        for (int i = 0; i < d; i++) {
            sum += q_ptr[q_offset + i] * k_ptr[k_offset + i];
        }
        
        attn_ptr[idx] = sum;
    }
}

// Attention-value multiplication: attn @ V
__global__ void av_kernel(
    const float* __restrict__ attn_ptr,
    const float* __restrict__ v_ptr,
    float* __restrict__ out_ptr,
    int B, int H, int T_q, int T_k, int d
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * H * T_q * d;
    
    if (idx < total) {
        int b = idx / (H * T_q * d);
        int h = (idx / (T_q * d)) % H;
        int q_idx = (idx / d) % T_q;
        int i = idx % d;
        
        int attn_offset = b * H * T_q * T_k + h * T_q * T_k + q_idx * T_k;
        
        float sum = 0.0f;
        for (int kv_idx = 0; kv_idx < T_k; kv_idx++) {
            int v_offset = b * H * T_k * d + h * T_k * d + kv_idx * d + i;
            sum += attn_ptr[attn_offset + kv_idx] * v_ptr[v_offset];
        }
        
        out_ptr[idx] = sum;
    }
}

// Host functions
void rope_cuda(
    const torch::Tensor& x,
    const torch::Tensor& cos,
    const torch::Tensor& sin,
    torch::Tensor& out
) {
    int B = x.size(0);
    int H = x.size(1);
    int L = x.size(2);
    int d = x.size(3);
    
    int total = B * H * L * (d / 2);
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    
    rope_kernel<<<blocks, threads>>>(
        x.data_ptr<float>(),
        cos.data_ptr<float>(),
        sin.data_ptr<float>(),
        out.data_ptr<float>(),
        B, H, L, d
    );
}

void softmax_cuda(
    torch::Tensor& attn,
    int B, int H, int T_q, int T_k,
    float scale
) {
    int total = B * H * T_q;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    
    softmax_kernel<<<blocks, threads>>>(
        attn.data_ptr<float>(),
        B, H, T_q, T_k, scale
    );
}

void qk_cuda(
    const torch::Tensor& q,
    const torch::Tensor& k,
    torch::Tensor& attn,
    int B, int H, int T_q, int T_k, int d
) {
    int total = B * H * T_q * T_k;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    
    qk_kernel<<<blocks, threads>>>(
        q.data_ptr<float>(),
        k.data_ptr<float>(),
        attn.data_ptr<float>(),
        B, H, T_q, T_k, d
    );
}

void av_cuda(
    const torch::Tensor& attn,
    const torch::Tensor& v,
    torch::Tensor& out,
    int B, int H, int T_q, int T_k, int d
) {
    int total = B * H * T_q * d;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    
    av_kernel<<<blocks, threads>>>(
        attn.data_ptr<float>(),
        v.data_ptr<float>(),
        out.data_ptr<float>(),
        B, H, T_q, T_k, d
    );
}
"""

cpp_source = """
void rope_cuda(const torch::Tensor& x, const torch::Tensor& cos, const torch::Tensor& sin, torch::Tensor& out);
void softmax_cuda(torch::Tensor& attn, int B, int H, int T_q, int T_k, float scale);
void qk_cuda(const torch::Tensor& q, const torch::Tensor& k, torch::Tensor& attn, int B, int H, int T_q, int T_k, int d);
void av_cuda(const torch::Tensor& attn, const torch::Tensor& v, torch::Tensor& out, int B, int H, int T_q, int T_k, int d);
"""

# Load the extension
attention_ext = load_inline(
    name='attention_ext',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['rope_cuda', 'softmax_cuda', 'qk_cuda', 'av_cuda'],
    extra_cuda_cflags=['-O3'],
    verbose=False
)


class Model(nn.Module):
    """
    Single-layer multi-head cross-attention with RoPE on Q and K, attention
    dropout, and output projection—intended as a fusion target for robotics
    stacks (e.g. proprio / action queries attending to vision or memory).
    """

    def __init__(
        self,
        dim: int,
        num_heads: int,
        dropout: float = 0.0,
        rope_base: float = 10000.0,
    ):
        super().__init__()
        assert dim % num_heads == 0
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        assert self.head_dim % 2 == 0, "RoPE requires an even head dimension"

        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)
        self.out_proj = nn.Linear(dim, dim, bias=False)
        self.dropout = nn.Dropout(dropout)
        self.scale = 1.0 / math.sqrt(self.head_dim)

        inv_freq = 1.0 / (
            rope_base ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim)
        )
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    @staticmethod
    def _rotate_half(x: torch.Tensor) -> torch.Tensor:
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat([-x2, x1], dim=-1)

    def _rope(self, x: torch.Tensor, seq_len: int) -> torch.Tensor:
        """Apply RoPE along the sequence dimension (last-but-one after head). x: (B, H, L, d)."""
        t = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(t, self.inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        cos = emb.cos().view(1, 1, seq_len, self.head_dim)
        sin = emb.sin().view(1, 1, seq_len, self.head_dim)
        return x * cos + self._rotate_half(x) * sin

    def forward(self, x_query: torch.Tensor, x_kv: torch.Tensor) -> torch.Tensor:
        B, T_q, _ = x_query.shape
        _, T_k, _ = x_kv.shape
        h, d = self.num_heads, self.head_dim

        q = self.q_proj(x_query).view(B, T_q, h, d).transpose(1, 2)
        k = self.k_proj(x_kv).view(B, T_k, h, d).transpose(1, 2)
        v = self.v_proj(x_kv).view(B, T_k, h, d).transpose(1, 2)

        q = self._rope(q, T_q)
        k = self._rope(k, T_k)

        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
        attn = F.softmax(attn, dim=-1)
        attn = self.dropout(attn)
        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).contiguous().view(B, T_q, self.dim)
        return self.out_proj(out)


class ModelNew(nn.Module):
    """
    Custom CUDA implementation of multi-head cross-attention with RoPE.
    """

    def __init__(
        self,
        dim: int,
        num_heads: int,
        dropout: float = 0.0,
        rope_base: float = 10000.0,
    ):
        super().__init__()
        assert dim % num_heads == 0
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        assert self.head_dim % 2 == 0, "RoPE requires an even head dimension"

        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)
        self.out_proj = nn.Linear(dim, dim, bias=False)
        self.dropout = nn.Dropout(dropout)
        self.scale = 1.0 / math.sqrt(self.head_dim)

        inv_freq = 1.0 / (
            rope_base ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim)
        )
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def _compute_rope_emb(self, seq_len: int, device):
        """Compute RoPE embeddings for a sequence length."""
        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(t, self.inv_freq)
        cos = freqs.cos()  # (seq_len, head_dim/2)
        sin = freqs.sin()  # (seq_len, head_dim/2)
        return cos, sin

    def forward(self, x_query: torch.Tensor, x_kv: torch.Tensor) -> torch.Tensor:
        B, T_q, _ = x_query.shape
        _, T_k, _ = x_kv.shape
        h, d = self.num_heads, self.head_dim

        # Project and reshape
        q = self.q_proj(x_query).view(B, T_q, h, d).transpose(1, 2).contiguous()
        k = self.k_proj(x_kv).view(B, T_k, h, d).transpose(1, 2).contiguous()
        v = self.v_proj(x_kv).view(B, T_k, h, d).transpose(1, 2).contiguous()

        # Compute RoPE embeddings
        cos_q, sin_q = self._compute_rope_emb(T_q, x_query.device)
        cos_k, sin_k = self._compute_rope_emb(T_k, x_kv.device)

        # Apply RoPE using custom CUDA kernel
        q_rope = torch.empty_like(q)
        k_rope = torch.empty_like(k)
        attention_ext.rope_cuda(q, cos_q, sin_q, q_rope)
        attention_ext.rope_cuda(k, cos_k, sin_k, k_rope)

        # Compute attention scores using custom CUDA kernel
        attn = torch.empty(B, h, T_q, T_k, device=x_query.device, dtype=x_query.dtype)
        attention_ext.qk_cuda(q_rope, k_rope, attn, B, h, T_q, T_k, d)

        # Apply softmax using custom CUDA kernel
        attention_ext.softmax_cuda(attn, B, h, T_q, T_k, self.scale)

        # Apply dropout (p=0 in test, so no-op)
        attn = self.dropout(attn)

        # Compute attention output using custom CUDA kernel
        attn_out = torch.empty(B, h, T_q, d, device=x_query.device, dtype=x_query.dtype)
        attention_ext.av_cuda(attn, v, attn_out, B, h, T_q, T_k, d)

        # Transpose and reshape for output projection
        out = attn_out.transpose(1, 2).contiguous().view(B, T_q, self.dim)

        return self.out_proj(out)


batch_size = 4
seq_len_q = 64
seq_len_kv = 128
dim = 256
num_heads = 8
dropout_p = 0.0


def get_inputs():
    x_query = torch.randn(batch_size, seq_len_q, dim)
    x_kv = torch.randn(batch_size, seq_len_kv, dim)
    return [x_query, x_kv]


def get_init_inputs():
    return [dim, num_heads, dropout_p]