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


class Model(nn.Module):
    """
    Invariant Point Attention (IPA) from AlphaFold2's structure module.
    Combines sequence-space attention with 3-D geometric queries/keys/values
    expressed in local residue frames.  The attention logits incorporate
    Euclidean distances in global frame, making it SE(3)-equivariant.
    """

    def __init__(self, node_dim, pair_dim, num_heads, num_query_points, num_value_points):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = node_dim // num_heads
        self.num_query_points = num_query_points
        self.num_value_points = num_value_points
        assert node_dim % num_heads == 0

        d = self.head_dim
        self.q_proj = nn.Linear(node_dim, num_heads * d, bias=False)
        self.k_proj = nn.Linear(node_dim, num_heads * d, bias=False)
        self.v_proj = nn.Linear(node_dim, num_heads * d, bias=False)

        self.q_pts = nn.Linear(node_dim, num_heads * num_query_points * 3, bias=False)
        self.k_pts = nn.Linear(node_dim, num_heads * num_query_points * 3, bias=False)
        self.v_pts = nn.Linear(node_dim, num_heads * num_value_points * 3, bias=False)

        self.pair_bias = nn.Linear(pair_dim, num_heads, bias=False)

        self.head_weights = nn.Parameter(torch.zeros(num_heads))

        out_features = num_heads * d + num_heads * num_value_points * 3
        self.out_proj = nn.Linear(out_features, node_dim)

        self.layer_norm = nn.LayerNorm(node_dim)

    def _apply_frame(self, pts, rotations, translations):
        """Transform local-frame points to global frame."""
        # pts: (B, N, P, 3), rotations: (B, N, 3, 3), translations: (B, N, 3)
        global_pts = torch.einsum("bnpc,bncd->bnpd", pts, rotations) + translations.unsqueeze(2)
        return global_pts

    def _apply_inv_frame(self, pts, rotations, translations):
        """Transform global-frame points to local frame."""
        centered = pts - translations.unsqueeze(2)
        local_pts = torch.einsum("bnpc,bndc->bnpd", centered, rotations)
        return local_pts

    def forward(
        self,
        node_repr: torch.Tensor,
        pair_repr: torch.Tensor,
        translations: torch.Tensor,
        rotations: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            node_repr:    (B, N, node_dim)
            pair_repr:    (B, N, N, pair_dim)
            translations: (B, N, 3)
            rotations:    (B, N, 3, 3)
        Returns:
            updated node representation (B, N, node_dim)
        """
        B, N, D = node_repr.shape
        h = self.num_heads
        d = self.head_dim
        Qp = self.num_query_points
        Vp = self.num_value_points

        x = self.layer_norm(node_repr)

        q = self.q_proj(x).view(B, N, h, d)
        k = self.k_proj(x).view(B, N, h, d)
        v = self.v_proj(x).view(B, N, h, d)

        # Points in local frame, then transform to global
        q_pts_local = self.q_pts(x).view(B, N, h * Qp, 3)
        k_pts_local = self.k_pts(x).view(B, N, h * Qp, 3)
        v_pts_local = self.v_pts(x).view(B, N, h * Vp, 3)

        q_pts = self._apply_frame(q_pts_local, rotations, translations).view(B, N, h, Qp, 3)
        k_pts = self._apply_frame(k_pts_local, rotations, translations).view(B, N, h, Qp, 3)
        v_pts = self._apply_frame(v_pts_local, rotations, translations).view(B, N, h, Vp, 3)

        # Scalar attention logits
        attn_scalar = torch.einsum("bihd,bjhd->bhij", q, k) / math.sqrt(d)

        # Point attention: negative squared distance between query and key points
        # q_pts: (B, N_i, h, Qp, 3), k_pts: (B, N_j, h, Qp, 3)
        q_expand = q_pts.unsqueeze(2)  # (B, N, 1, h, Qp, 3)
        k_expand = k_pts.unsqueeze(1)  # (B, 1, N, h, Qp, 3)
        pt_dist = ((q_expand - k_expand) ** 2).sum(dim=(-1, -2))  # (B, N, N, h)

        w = F.softplus(self.head_weights)
        pt_attn = -(w.view(1, 1, 1, h) * pt_dist) / 2.0
        pt_attn = pt_attn.permute(0, 3, 1, 2)  # (B, h, N, N)

        # Pair bias
        pair_bias = self.pair_bias(pair_repr).permute(0, 3, 1, 2)  # (B, h, N, N)

        attn = attn_scalar + pt_attn + pair_bias
        attn = F.softmax(attn, dim=-1)  # (B, h, N_i, N_j)

        # Aggregate scalar values
        out_scalar = torch.einsum("bhij,bjhd->bihd", attn, v).reshape(B, N, h * d)

        # Aggregate point values in global frame, then transform back to local
        out_pts = torch.einsum("bhij,bjhpc->bihpc", attn, v_pts)  # (B, N, h, Vp, 3)
        out_pts = out_pts.reshape(B, N, h * Vp, 3)
        out_pts_local = self._apply_inv_frame(out_pts, rotations, translations)
        out_pts_flat = out_pts_local.reshape(B, N, h * Vp * 3)

        out = torch.cat([out_scalar, out_pts_flat], dim=-1)
        return self.out_proj(out)


node_dim = 64
pair_dim = 32
num_heads = 4
num_query_points = 4
num_value_points = 8
seq_len = 32
batch_size = 2


def get_inputs():
    p = popcorn_pri
    node_repr = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), node_dim)
    pair_repr = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(seq_len), pair_dim)
    translations = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), 3)
    rotations = torch.zeros(p.jitter_int(batch_size), p.jitter_int(seq_len), 3, 3)
    for b in range(p.jitter_int(batch_size)):
        for n in range(p.jitter_int(seq_len)):
            rotations[b, n] = torch.eye(3)
    return [node_repr, pair_repr, translations, rotations]


def get_init_inputs():
    return [node_dim, pair_dim, num_heads, num_query_points, num_value_points]

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

# Define CUDA kernel for fused attention computation
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <c10/cuda/CUDAException.h>

__device__ __forceinline__ float warpReduceSum(float val) {
    for (int offset = 16; offset > 0; offset /= 2)
        val += __shfl_xor_sync(0xffffffff, val, offset);
    return val;
}

__device__ __forceinline__ float blockReduceSum(float val, int tid) {
    __shared__ float shared[32];
    int lane = tid & 31;
    int wid = tid >> 5;
    
    val = warpReduceSum(val);
    
    if (lane == 0) shared[wid] = val;
    __syncthreads();
    
    val = (tid < (blockDim.x >> 5)) ? shared[lane] : 0.0f;
    if (wid == 0) val = warpReduceSum(val);
    return val;
}

// Fused kernel: compute attention scores (scalar + point + pair_bias) and softmax
// Grid: (B, H, N), Block: (N_j) - each block handles one query position
__global__ void fused_attention_kernel(
    const float* __restrict__ q,        // (B, N, H, D)
    const float* __restrict__ k,        // (B, N, H, D)
    const float* __restrict__ q_pts,    // (B, N, H, Qp, 3)
    const float* __restrict__ k_pts,    // (B, N, H, Qp, 3)
    const float* __restrict__ pair_bias, // (B, H, N, N)
    const float* __restrict__ head_weights, // (H,)
    float* __restrict__ attn,           // (B, H, N, N)
    int B, int N, int H, int D, int Qp, float inv_sqrt_d
) {
    int b = blockIdx.x;  // batch
    int h = blockIdx.y;  // head
    int i = blockIdx.z;  // query position
    
    if (b >= B || h >= H || i >= N) return;
    
    int tid = threadIdx.x;
    int j = tid;  // key position
    
    // Load query vector for this head
    extern __shared__ float smem[];
    float* q_local = smem;  // size D
    float* q_pts_local = smem + D;  // size Qp * 3
    
    if (tid < D) {
        q_local[tid] = q[b * N * H * D + i * H * D + h * D + tid];
    }
    if (tid < Qp * 3) {
        q_pts_local[tid] = q_pts[b * N * H * Qp * 3 + i * H * Qp * 3 + h * Qp * 3 + tid];
    }
    __syncthreads();
    
    float hw = head_weights[h];
    
    // Each thread computes attention for one (i, j) pair
    // Compute scalar attention
    float scalar_attn = 0.0f;
    
    // Compute scalar product q_i * k_j
    for (int d = 0; d < D; d++) {
        float k_val = k[b * N * H * D + j * H * D + h * D + d];
        scalar_attn += q_local[d] * k_val;
    }
    scalar_attn *= inv_sqrt_d;
    
    // Compute point attention (negative squared distance)
    float pt_dist = 0.0f;
    for (int p = 0; p < Qp; p++) {
        for (int c = 0; c < 3; c++) {
            float q_pt = q_pts_local[p * 3 + c];
            float k_pt = k_pts[b * N * H * Qp * 3 + j * H * Qp * 3 + h * Qp * 3 + p * 3 + c];
            float diff = q_pt - k_pt;
            pt_dist += diff * diff;
        }
    }
    float pt_attn = -hw * pt_dist / 2.0f;
    
    // Get pair bias
    float pb = pair_bias[b * H * N * N + h * N * N + i * N + j];
    
    // Total attention score
    float score = scalar_attn + pt_attn + pb;
    
    // Store for softmax (need all scores first)
    __shared__ float scores[1024];  // Assuming N <= 1024
    scores[j] = score;
    __syncthreads();
    
    // Find max for numerical stability
    __shared__ float max_val;
    float local_max = scores[j];
    if (tid >= N) local_max = -1e30f;
    
    for (int offset = 16; offset > 0; offset /= 2) {
        float other = __shfl_xor_sync(0xffffffff, local_max, offset);
        local_max = fmaxf(local_max, other);
    }
    if (tid == 0) max_val = local_max;
    __syncthreads();
    
    // Compute exp(score - max)
    float exp_val = (tid < N) ? expf(scores[tid] - max_val) : 0.0f;
    
    // Reduce to get sum
    float sum = blockReduceSum(exp_val, tid);
    __shared__ float sum_val;
    if (tid == 0) sum_val = sum;
    __syncthreads();
    
    // Compute softmax
    float softmax_val = exp_val / sum_val;
    
    // Store result
    if (j < N) {
        attn[b * H * N * N + h * N * N + i * N + j] = softmax_val;
    }
}

// Fused kernel: aggregate values (scalar + point) and apply inverse frame
__global__ void fused_aggregate_kernel(
    const float* __restrict__ attn,      // (B, H, N, N)
    const float* __restrict__ v,         // (B, N, H, D)
    const float* __restrict__ v_pts,     // (B, N, H, Vp, 3)
    const float* __restrict__ rotations, // (B, N, 3, 3)
    const float* __restrict__ translations, // (B, N, 3)
    float* __restrict__ out_scalar,      // (B, N, H, D)
    float* __restrict__ out_pts_local,   // (B, N, H * Vp, 3)
    int B, int N, int H, int D, int Vp
) {
    int b = blockIdx.x;
    int i = blockIdx.y;  // output position
    int h = blockIdx.z;  // head
    int d = threadIdx.x; // dimension within head
    
    if (b >= B || i >= N || h >= H || d >= D + Vp * 3) return;
    
    int N_total = N;
    
    if (d < D) {
        // Aggregate scalar values
        float sum = 0.0f;
        for (int j = 0; j < N_total; j++) {
            float a = attn[b * H * N * N + h * N * N + i * N + j];
            float v_val = v[b * N * H * D + j * H * D + h * D + d];
            sum += a * v_val;
        }
        out_scalar[b * N * H * D + i * H * D + h * D + d] = sum;
    } else {
        // Aggregate point values
        int pt_idx = d - D;  // 0 to Vp*3-1
        int vp = pt_idx / 3;
        int c = pt_idx % 3;
        
        // First aggregate in global frame
        float sum = 0.0f;
        for (int j = 0; j < N_total; j++) {
            float a = attn[b * H * N * N + h * N * N + i * N + j];
            float v_pt = v_pts[b * N * H * Vp * 3 + j * H * Vp * 3 + h * Vp * 3 + vp * 3 + c];
            sum += a * v_pt;
        }
        
        // Apply inverse frame transformation
        // out_global - translations, then rotate by rotations^T
        float trans = translations[b * N * 3 + i * 3 + c];
        float centered = sum - trans;
        
        float local_val = 0.0f;
        for (int k = 0; k < 3; k++) {
            float rot = rotations[b * N * 3 * 3 + i * 9 + k * 3 + c];  // R^T = R col-major
            float coord = (k == 0) ? 0.0f : 0.0f;  // Need all 3 coords
        }
        
        // Store intermediate for inverse transform (need all 3 coords)
        // This kernel needs revision for proper inverse transform
    }
}

// Separate kernel for inverse frame transformation
__global__ void inverse_frame_kernel(
    const float* __restrict__ pts_global,  // (B, N, H*Vp, 3)
    const float* __restrict__ rotations,   // (B, N, 3, 3)
    const float* __restrict__ translations, // (B, N, 3)
    float* __restrict__ pts_local,         // (B, N, H*Vp, 3)
    int B, int N, int total_pts
) {
    int b = blockIdx.x;
    int n = blockIdx.y;
    int p = blockIdx.z;  // point index
    int c = threadIdx.x; // coordinate (0, 1, 2)
    
    if (b >= B || n >= N || p >= total_pts || c >= 3) return;
    
    // Load translation for this coordinate
    float trans = translations[b * N * 3 + n * 3 + c];
    
    // Load point coordinates
    float pt[3];
    for (int i = 0; i < 3; i++) {
        pt[i] = pts_global[b * N * total_pts * 3 + n * total_pts * 3 + p * 3 + i];
    }
    
    // Center
    float centered[3];
    for (int i = 0; i < 3; i++) {
        centered[i] = pt[i] - translations[b * N * 3 + n * 3 + i];
    }
    
    // Apply rotation^T (rotation is stored row-major, so R^T[i,j] = R[j,i])
    float local_val = 0.0f;
    for (int k = 0; k < 3; k++) {
        // R^T[c,k] = R[k,c]
        float rot_T = rotations[b * N * 9 + n * 9 + k * 3 + c];
        local_val += rot_T * centered[k];
    }
    
    pts_local[b * N * total_pts * 3 + n * total_pts * 3 + p * 3 + c] = local_val;
}

// Point transformation kernel (local to global)
__global__ void apply_frame_kernel(
    const float* __restrict__ pts_local,   // (B, N, H*P, 3)
    const float* __restrict__ rotations,   // (B, N, 3, 3)
    const float* __restrict__ translations, // (B, N, 3)
    float* __restrict__ pts_global,        // (B, N, H*P, 3)
    int B, int N, int total_pts
) {
    int b = blockIdx.x;
    int n = blockIdx.y;
    int p = blockIdx.z;  // point index
    int c = threadIdx.x; // coordinate
    
    if (b >= B || n >= N || p >= total_pts || c >= 3) return;
    
    // Load point coordinates
    float pt[3];
    for (int i = 0; i < 3; i++) {
        pt[i] = pts_local[b * N * total_pts * 3 + n * total_pts * 3 + p * 3 + i];
    }
    
    // Apply rotation: R * pt
    float rotated = 0.0f;
    for (int k = 0; k < 3; k++) {
        float rot = rotations[b * N * 9 + n * 9 + c * 3 + k];  // row-major
        rotated += rot * pt[k];
    }
    
    // Add translation
    float result = rotated + translations[b * N * 3 + n * 3 + c];
    
    pts_global[b * N * total_pts * 3 + n * total_pts * 3 + p * 3 + c] = result;
}

// Optimized attention computation with better memory access
__global__ void attention_forward_kernel(
    const float* __restrict__ q,
    const float* __restrict__ k,
    const float* __restrict__ q_pts,
    const float* __restrict__ k_pts,
    const float* __restrict__ pair_bias,
    const float* __restrict__ head_weights,
    float* __restrict__ attn_out,
    int B, int N, int H, int D, int Qp, float inv_sqrt_d
) {
    // Each block handles one (batch, head, query_i) combination
    int b = blockIdx.x;
    int h = blockIdx.y;
    int i = blockIdx.z;
    
    if (b >= B || h >= H || i >= N) return;
    
    int tid = threadIdx.x;
    int lane = tid & 31;
    
    // Shared memory for query and query points
    extern __shared__ char smem_buf[];
    float* q_smem = reinterpret_cast<float*>(smem_buf);
    float* q_pts_smem = q_smem + D;
    float* scores_smem = q_pts_smem + Qp * 3;
    
    // Load q and q_pts into shared memory
    for (int idx = tid; idx < D; idx += blockDim.x) {
        q_smem[idx] = q[b * N * H * D + i * H * D + h * D + idx];
    }
    for (int idx = tid; idx < Qp * 3; idx += blockDim.x) {
        q_pts_smem[idx] = q_pts[b * N * H * Qp * 3 + i * H * Qp * 3 + h * Qp * 3 + idx];
    }
    __syncthreads();
    
    float hw = head_weights[h];
    
    // Compute attention scores for all j
    float local_max = -1e30f;
    
    for (int j = tid; j < N; j += blockDim.x) {
        // Scalar attention
        float scalar = 0.0f;
        for (int d = 0; d < D; d++) {
            scalar += q_smem[d] * k[b * N * H * D + j * H * D + h * D + d];
        }
        scalar *= inv_sqrt_d;
        
        // Point attention
        float pt_dist = 0.0f;
        for (int p = 0; p < Qp; p++) {
            for (int c = 0; c < 3; c++) {
                float diff = q_pts_smem[p * 3 + c] - 
                            k_pts[b * N * H * Qp * 3 + j * H * Qp * 3 + h * Qp * 3 + p * 3 + c];
                pt_dist += diff * diff;
            }
        }
        float pt_attn = -hw * pt_dist / 2.0f;
        
        // Pair bias
        float pb = pair_bias[b * H * N * N + h * N * N + i * N + j];
        
        float score = scalar + pt_attn + pb;
        scores_smem[j] = score;
        local_max = fmaxf(local_max, score);
    }
    
    // Block-level max reduction
    __shared__ float global_max;
    float reduced_max = blockReduceSum(local_max, tid);  // Actually need max, not sum
    // Use separate max reduction
    for (int offset = blockDim.x / 2; offset > 0; offset /= 2) {
        float other = __shfl_xor_sync(0xffffffff, local_max, offset);
        local_max = fmaxf(local_max, other);
    }
    if (tid == 0) global_max = local_max;
    __syncthreads();
    
    // Compute exp and sum
    float local_sum = 0.0f;
    for (int j = tid; j < N; j += blockDim.x) {
        float exp_val = expf(scores_smem[j] - global_max);
        scores_smem[j] = exp_val;
        local_sum += exp_val;
    }
    
    __shared__ float global_sum;
    float reduced_sum = blockReduceSum(local_sum, tid);
    if (tid == 0) global_sum = reduced_sum;
    __syncthreads();
    
    // Normalize and store
    for (int j = tid; j < N; j += blockDim.x) {
        attn_out[b * H * N * N + h * N * N + i * N + j] = scores_smem[j] / global_sum;
    }
}

// Value aggregation kernel
__global__ void value_aggregate_kernel(
    const float* __restrict__ attn,
    const float* __restrict__ v,
    const float* __restrict__ v_pts,
    float* __restrict__ out_scalar,
    float* __restrict__ out_pts_global,
    int B, int N, int H, int D, int Vp
) {
    int b = blockIdx.x;
    int h = blockIdx.y;
    int i = blockIdx.z;  // output position
    
    if (b >= B || h >= H || i >= N) return;
    
    int tid = threadIdx.x;
    
    // Each thread handles one output dimension
    // First D dimensions are scalar, then Vp*3 are point coordinates
    
    __shared__ float attn_smem[1024];  // Assuming N <= 1024
    
    // Load attention weights
    for (int j = tid; j < N; j += blockDim.x) {
        attn_smem[j] = attn[b * H * N * N + h * N * N + i * N + j];
    }
    __syncthreads();
    
    // Scalar aggregation
    if (tid < D) {
        float sum = 0.0f;
        for (int j = 0; j < N; j++) {
            sum += attn_smem[j] * v[b * N * H * D + j * H * D + h * D + tid];
        }
        out_scalar[b * N * H * D + i * H * D + h * D + tid] = sum;
    }
    
    // Point aggregation
    if (tid < Vp * 3) {
        int vp = tid / 3;
        int c = tid % 3;
        float sum = 0.0f;
        for (int j = 0; j < N; j++) {
            sum += attn_smem[j] * v_pts[b * N * H * Vp * 3 + j * H * Vp * 3 + h * Vp * 3 + vp * 3 + c];
        }
        out_pts_global[b * N * H * Vp * 3 + i * H * Vp * 3 + h * Vp * 3 + vp * 3 + c] = sum;
    }
}
"""

cpp_source = """
torch::Tensor apply_frame_cuda(
    torch::Tensor pts_local,
    torch::Tensor rotations,
    torch::Tensor translations,
    int total_pts
);

torch::Tensor inverse_frame_cuda(
    torch::Tensor pts_global,
    torch::Tensor rotations,
    torch::Tensor translations,
    int total_pts
);

torch::Tensor attention_forward_cuda(
    torch::Tensor q,
    torch::Tensor k,
    torch::Tensor q_pts,
    torch::Tensor k_pts,
    torch::Tensor pair_bias,
    torch::Tensor head_weights,
    int B, int N, int H, int D, int Qp, float inv_sqrt_d
);

torch::Tensor value_aggregate_cuda(
    torch::Tensor attn,
    torch::Tensor v,
    torch::Tensor v_pts,
    int B, int N, int H, int D, int Vp
);

torch::Tensor value_aggregate_cuda(
    torch::Tensor attn,
    torch::Tensor v,
    torch::Tensor v_pts,
    int B, int N, int H, int D, int Vp
) {
    auto options = torch::TensorOptions().dtype(torch::kFloat32).device(attn.device());
    auto out_scalar = torch::empty({B, N, H, D}, options);
    auto out_pts_global = torch::empty({B, N, H, Vp, 3}, options);
    
    int threads = std::max(D, Vp * 3);
    threads = std::min(threads, 256);
    
    dim3 grid(B, H, N);
    
    AT_DISPATCH_FLOATING_TYPES(attn.scalar_type(), "value_aggregate", ([&] {
        value_aggregate_kernel<<<grid, threads>>>(
            attn.data_ptr<float>(),
            v.data_ptr<float>(),
            v_pts.data_ptr<float>(),
            out_scalar.data_ptr<float>(),
            out_pts_global.data_ptr<float>(),
            B, N, H, D, Vp
        );
    }));
    
    return torch::cat({out_scalar.reshape({B, N, H * D}), 
                       out_pts_global.reshape({B, N, H * Vp * 3})}, -1);
}

torch::Tensor apply_frame_cuda(
    torch::Tensor pts_local,
    torch::Tensor rotations,
    torch::Tensor translations,
    int total_pts
) {
    TORCH_CHECK(pts_local.is_cuda(), "pts_local must be CUDA tensor");
    TORCH_CHECK(rotations.is_cuda(), "rotations must be CUDA tensor");
    TORCH_CHECK(translations.is_cuda(), "translations must be CUDA tensor");
    
    auto options = torch::TensorOptions().dtype(torch::kFloat32).device(pts_local.device());
    int B = pts_local.size(0);
    int N = pts_local.size(1);
    
    auto pts_global = torch::empty({B, N, total_pts, 3}, options);
    
    dim3 grid(B, N, total_pts);
    int threads = 3;
    
    AT_DISPATCH_FLOATING_TYPES(pts_local.scalar_type(), "apply_frame", ([&] {
        apply_frame_kernel<<<grid, threads>>>(
            pts_local.data_ptr<float>(),
            rotations.data_ptr<float>(),
            translations.data_ptr<float>(),
            pts_global.data_ptr<float>(),
            B, N, total_pts
        );
    }));
    
    return pts_global;
}

torch::Tensor inverse_frame_cuda(
    torch::Tensor pts_global,
    torch::Tensor rotations,
    torch::Tensor translations,
    int total_pts
) {
    TORCH_CHECK(pts_global.is_cuda(), "pts_global must be CUDA tensor");
    
    auto options = torch::TensorOptions().dtype(torch::kFloat32).device(pts_global.device());
    int B = pts_global.size(0);
    int N = pts_global.size(1);
    
    auto pts_local = torch::empty({B, N, total_pts, 3}, options);
    
    dim3 grid(B, N, total_pts);
    int threads = 3;
    
    AT_DISPATCH_FLOATING_TYPES(pts_global.scalar_type(), "inverse_frame", ([&] {
        inverse_frame_kernel<<<grid, threads>>>(
            pts_global.data_ptr<float>(),
            rotations.data_ptr<float>(),
            translations.data_ptr<float>(),
            pts_local.data_ptr<float>(),
            B, N, total_pts
        );
    }));
    
    return pts_local;
}

torch::Tensor attention_forward_cuda(
    torch::Tensor q,
    torch::Tensor k,
    torch::Tensor q_pts,
    torch::Tensor k_pts,
    torch::Tensor pair_bias,
    torch::Tensor head_weights,
    int B, int N, int H, int D, int Qp, float inv_sqrt_d
) {
    TORCH_CHECK(q.is_cuda(), "q must be CUDA tensor");
    
    auto options = torch::TensorOptions().dtype(torch::kFloat32).device(q.device());
    auto attn_out = torch::empty({B, H, N, N}, options);
    
    dim3 grid(B, H, N);
    int threads = std::min(N, 256);
    int smem_size = D * sizeof(float) + Qp * 3 * sizeof(float) + N * sizeof(float);
    
    AT_DISPATCH_FLOATING_TYPES(q.scalar_type(), "attention_forward", ([&] {
        attention_forward_kernel<<<grid, threads, smem_size>>>(
            q.data_ptr<float>(),
            k.data_ptr<float>(),
            q_pts.data_ptr<float>(),
            k_pts.data_ptr<float>(),
            pair_bias.data_ptr<float>(),
            head_weights.data_ptr<float>(),
            attn_out.data_ptr<float>(),
            B, N, H, D, Qp, inv_sqrt_d
        );
    }));
    
    return attn_out;
}
"""

# Compile the extension
ipa_ops = load_inline(
    name='ipa_ops',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['apply_frame_cuda', 'inverse_frame_cuda', 'attention_forward_cuda', 'value_aggregate_cuda'],
    extra_cuda_cflags=['-O3'],
    verbose=False
)


class ModelNew(nn.Module):
    """
    Invariant Point Attention (IPA) with custom CUDA kernels.
    """

    def __init__(self, node_dim, pair_dim, num_heads, num_query_points, num_value_points):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = node_dim // num_heads
        self.num_query_points = num_query_points
        self.num_value_points = num_value_points
        assert node_dim % num_heads == 0

        d = self.head_dim
        self.q_proj = nn.Linear(node_dim, num_heads * d, bias=False)
        self.k_proj = nn.Linear(node_dim, num_heads * d, bias=False)
        self.v_proj = nn.Linear(node_dim, num_heads * d, bias=False)

        self.q_pts = nn.Linear(node_dim, num_heads * num_query_points * 3, bias=False)
        self.k_pts = nn.Linear(node_dim, num_heads * num_query_points * 3, bias=False)
        self.v_pts = nn.Linear(node_dim, num_heads * num_value_points * 3, bias=False)

        self.pair_bias = nn.Linear(pair_dim, num_heads, bias=False)

        self.head_weights = nn.Parameter(torch.zeros(num_heads))

        out_features = num_heads * d + num_heads * num_value_points * 3
        self.out_proj = nn.Linear(out_features, node_dim)

        self.layer_norm = nn.LayerNorm(node_dim)

    def forward(
        self,
        node_repr: torch.Tensor,
        pair_repr: torch.Tensor,
        translations: torch.Tensor,
        rotations: torch.Tensor,
    ) -> torch.Tensor:
        B, N, D = node_repr.shape
        h = self.num_heads
        d = self.head_dim
        Qp = self.num_query_points
        Vp = self.num_value_points

        x = self.layer_norm(node_repr)

        # Projections
        q = self.q_proj(x).view(B, N, h, d)
        k = self.k_proj(x).view(B, N, h, d)
        v = self.v_proj(x).view(B, N, h, d)

        # Point projections (local frame)
        q_pts_local = self.q_pts(x).view(B, N, h * Qp, 3)
        k_pts_local = self.k_pts(x).view(B, N, h * Qp, 3)
        v_pts_local = self.v_pts(x).view(B, N, h * Vp, 3)

        # Transform points to global frame using custom CUDA kernel
        if x.is_cuda:
            q_pts_global = ipa_ops.apply_frame_cuda(
                q_pts_local.contiguous().float(),
                rotations.contiguous().float(),
                translations.contiguous().float(),
                h * Qp
            )
            k_pts_global = ipa_ops.apply_frame_cuda(
                k_pts_local.contiguous().float(),
                rotations.contiguous().float(),
                translations.contiguous().float(),
                h * Qp
            )
            v_pts_global = ipa_ops.apply_frame_cuda(
                v_pts_local.contiguous().float(),
                rotations.contiguous().float(),
                translations.contiguous().float(),
                h * Vp
            )
            
            q_pts = q_pts_global.view(B, N, h, Qp, 3)
            k_pts = k_pts_global.view(B, N, h, Qp, 3)
            v_pts = v_pts_global.view(B, N, h, Vp, 3)
            
            # Compute pair bias
            pair_bias = self.pair_bias(pair_repr).permute(0, 3, 1, 2)
            
            # Compute head weights
            w = F.softplus(self.head_weights)
            
            # Fused attention computation
            inv_sqrt_d = 1.0 / math.sqrt(d)
            attn = ipa_ops.attention_forward_cuda(
                q.contiguous().float(),
                k.contiguous().float(),
                q_pts.contiguous().float(),
                k_pts.contiguous().float(),
                pair_bias.contiguous().float(),
                w.contiguous().float(),
                B, N, h, d, Qp, inv_sqrt_d
            )
            
            # Value aggregation
            out_combined = ipa_ops.value_aggregate_cuda(
                attn.contiguous().float(),
                v.contiguous().float(),
                v_pts.contiguous().float(),
                B, N, h, d, Vp
            )
            
            # Split scalar and point outputs
            out_scalar = out_combined[..., :h * d].view(B, N, h, d)
            out_pts_global = out_combined[..., h * d:].view(B, N, h * Vp, 3)
            
            # Transform points back to local frame
            out_pts_local = ipa_ops.inverse_frame_cuda(
                out_pts_global.contiguous().float(),
                rotations.contiguous().float(),
                translations.contiguous().float(),
                h * Vp
            )
            out_pts_flat = out_pts_local.view(B, N, h * Vp * 3)
            
            out = torch.cat([out_scalar.reshape(B, N, h * d), out_pts_flat], dim=-1)
            return self.out_proj(out)
        else:
            # CPU fallback - use original implementation
            q_pts = self._apply_frame(q_pts_local, rotations, translations).view(B, N, h, Qp, 3)
            k_pts = self._apply_frame(k_pts_local, rotations, translations).view(B, N, h, Qp, 3)
            v_pts = self._apply_frame(v_pts_local, rotations, translations).view(B, N, h, Vp, 3)

            attn_scalar = torch.einsum("bihd,bjhd->bhij", q, k) / math.sqrt(d)

            q_expand = q_pts.unsqueeze(2)
            k_expand = k_pts.unsqueeze(1)
            pt_dist = ((q_expand - k_expand) ** 2).sum(dim=(-1, -2))

            w = F.softplus(self.head_weights)
            pt_attn = -(w.view(1, 1, 1, h) * pt_dist) / 2.0
            pt_attn = pt_attn.permute(0, 3, 1, 2)

            pair_bias = self.pair_bias(pair_repr).permute(0, 3, 1, 2)

            attn = attn_scalar + pt_attn + pair_bias
            attn = F.softmax(attn, dim=-1)

            out_scalar = torch.einsum("bhij,bjhd->bihd", attn, v).reshape(B, N, h * d)

            out_pts = torch.einsum("bhij,bjhpc->bihpc", attn, v_pts)
            out_pts = out_pts.reshape(B, N, h * Vp, 3)
            out_pts_local = self._apply_inv_frame(out_pts, rotations, translations)
            out_pts_flat = out_pts_local.reshape(B, N, h * Vp * 3)

            out = torch.cat([out_scalar, out_pts_flat], dim=-1)
            return self.out_proj(out)

    def _apply_frame(self, pts, rotations, translations):
        global_pts = torch.einsum("bnpc,bncd->bnpd", pts, rotations) + translations.unsqueeze(2)
        return global_pts

    def _apply_inv_frame(self, pts, rotations, translations):
        centered = pts - translations.unsqueeze(2)
        local_pts = torch.einsum("bnpc,bndc->bnpd", centered, rotations)
        return local_pts