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

# Fused CUDA kernel for IPA attention and aggregation - optimized version
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <c10/cuda/CUDAException.h>

__device__ __forceinline__ float warp_reduce_max(float val) {
    for (int offset = 16; offset > 0; offset /= 2) {
        val = max(val, __shfl_down_sync(0xffffffff, val, offset));
    }
    return val;
}

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

// Optimized IPA attention kernel
// Each thread block handles one (batch, head, query_pos) tuple
__global__ void ipa_attention_kernel(
    const float* __restrict__ q,           // (B, N, h, d)
    const float* __restrict__ k,           // (B, N, h, d)
    const float* __restrict__ v,           // (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__ v_pts,       // (B, N, h, Vp, 3)
    const float* __restrict__ pair_bias,   // (B, N, N, h)
    const float* __restrict__ head_weights, // (h,)
    const float* __restrict__ rotations,    // (B, N, 9)
    const float* __restrict__ translations, // (B, N, 3)
    float* __restrict__ out_scalar,        // (B, N, h, d)
    float* __restrict__ out_pts,           // (B, N, h*Vp, 3)
    int B, int N, int h, int d, int Qp, int Vp
) {
    const int batch = blockIdx.y;
    const int head = blockIdx.z;
    const int query_pos = blockIdx.x;
    
    const int tid = threadIdx.x;
    const int warp_id = tid / 32;
    const int lane_id = tid % 32;
    const int block_size = blockDim.x;
    const int num_warps = block_size / 32;
    
    // Load head weight and apply softplus
    float w = head_weights[head];
    w = log1pf(expf(w));
    
    // Shared memory layout
    extern __shared__ char smem_base[];
    float* smem = reinterpret_cast<float*>(smem_base);
    float* q_local = smem;                         // d elements
    float* q_pts_local = smem + d;                 // Qp * 3 elements
    float* rot_local = smem + d + Qp * 3;          // 9 elements
    float* trans_local = smem + d + Qp * 3 + 9;    // 3 elements
    float* scores = smem + d + Qp * 3 + 9 + 3;     // N elements
    float* shared_pts = smem + d + Qp * 3 + 9 + 3 + N; // Vp * 3 elements
    
    // Base indices for current batch/head/query
    const int q_base = batch * N * h * d + query_pos * h * d + head * d;
    const int q_pts_base = batch * N * h * Qp * 3 + query_pos * h * Qp * 3 + head * Qp * 3;
    const int rot_base = batch * N * 9 + query_pos * 9;
    const int trans_base = batch * N * 3 + query_pos * 3;
    const int pair_bias_base = batch * N * N * h + query_pos * N * h + head;
    
    // Cooperatively load q, q_pts, rotation, translation
    for (int i = tid; i < d; i += block_size) {
        q_local[i] = q[q_base + i];
    }
    for (int i = tid; i < Qp * 3; i += block_size) {
        q_pts_local[i] = q_pts[q_pts_base + i];
    }
    for (int i = tid; i < 9; i += block_size) {
        rot_local[i] = rotations[rot_base + i];
    }
    for (int i = tid; i < 3; i += block_size) {
        trans_local[i] = translations[trans_base + i];
    }
    __syncthreads();
    
    // Precompute 1/sqrt(d)
    const float inv_sqrt_d = 1.0f / sqrtf((float)d);
    
    // Compute attention scores and find max
    float max_score = -INFINITY;
    for (int key_pos = tid; key_pos < N; key_pos += block_size) {
        // Scalar attention
        float scalar_score = 0.0f;
        const int k_base = batch * N * h * d + key_pos * h * d + head * d;
        for (int i = 0; i < d; i++) {
            scalar_score += q_local[i] * k[k_base + i];
        }
        scalar_score *= inv_sqrt_d;
        
        // Point attention
        float pt_score = 0.0f;
        const int k_pts_base = batch * N * h * Qp * 3 + key_pos * h * Qp * 3 + head * Qp * 3;
        for (int p = 0; p < Qp; p++) {
            for (int c = 0; c < 3; c++) {
                float diff = q_pts_local[p * 3 + c] - k_pts[k_pts_base + p * 3 + c];
                pt_score += diff * diff;
            }
        }
        pt_score = -w * pt_score * 0.5f;
        
        // Pair bias
        float pb = pair_bias[pair_bias_base + key_pos * h];
        
        float total_score = scalar_score + pt_score + pb;
        scores[key_pos] = total_score;
        max_score = max(max_score, total_score);
    }
    
    // Block reduce max
    __shared__ float block_max[32];
    float my_max = warp_reduce_max(max_score);
    if (lane_id == 0) {
        block_max[warp_id] = my_max;
    }
    __syncthreads();
    if (warp_id == 0) {
        my_max = lane_id < num_warps ? block_max[lane_id] : -INFINITY;
        my_max = warp_reduce_max(my_max);
        if (lane_id == 0) {
            block_max[0] = my_max;
        }
    }
    __syncthreads();
    float global_max = block_max[0];
    
    // Compute softmax
    float sum_exp = 0.0f;
    for (int i = tid; i < N; i += block_size) {
        float e = expf(scores[i] - global_max);
        scores[i] = e;
        sum_exp += e;
    }
    
    // Block reduce sum
    __shared__ float block_sum[32];
    float my_sum = warp_reduce_sum(sum_exp);
    if (lane_id == 0) {
        block_sum[warp_id] = my_sum;
    }
    __syncthreads();
    if (warp_id == 0) {
        my_sum = lane_id < num_warps ? block_sum[lane_id] : 0.0f;
        my_sum = warp_reduce_sum(my_sum);
        if (lane_id == 0) {
            block_sum[0] = my_sum;
        }
    }
    __syncthreads();
    float inv_sum = 1.0f / block_sum[0];
    
    // Normalize and compute weighted values in single pass
    // Each thread accumulates its portion
    float out_s[16] = {0};
    float out_p[24] = {0};
    
    for (int key_pos = tid; key_pos < N; key_pos += block_size) {
        float attn = scores[key_pos] * inv_sum;
        
        // Scalar value aggregation
        const int v_base = batch * N * h * d + key_pos * h * d + head * d;
        for (int i = 0; i < d; i++) {
            out_s[i] += attn * v[v_base + i];
        }
        
        // Point value aggregation
        const int v_pts_base = batch * N * h * Vp * 3 + key_pos * h * Vp * 3 + head * Vp * 3;
        for (int p = 0; p < Vp; p++) {
            for (int c = 0; c < 3; c++) {
                out_p[p * 3 + c] += attn * v_pts[v_pts_base + p * 3 + c];
            }
        }
    }
    
    // Warp reduce for scalar output and write directly
    for (int i = 0; i < d; i++) {
        float val = warp_reduce_sum(out_s[i]);
        if (lane_id == 0) {
            atomicAdd(&out_scalar[q_base + i], val);
        }
    }
    
    // Initialize shared_pts for point output
    for (int i = tid; i < Vp * 3; i += block_size) {
        shared_pts[i] = 0.0f;
    }
    __syncthreads();
    
    // Warp reduce for point output
    for (int p = 0; p < Vp; p++) {
        for (int c = 0; c < 3; c++) {
            float val = warp_reduce_sum(out_p[p * 3 + c]);
            if (lane_id == 0) {
                atomicAdd(&shared_pts[p * 3 + c], val);
            }
        }
    }
    __syncthreads();
    
    // Transform points to local frame and write output
    const int out_pts_base = batch * N * h * Vp * 3 + query_pos * h * Vp * 3 + head * Vp * 3;
    if (tid < Vp * 3) {
        int p = tid / 3;
        int c = tid % 3;
        float local_pt = 0.0f;
        for (int i = 0; i < 3; i++) {
            float global_pt = shared_pts[p * 3 + i] - trans_local[i];
            local_pt += global_pt * rot_local[c * 3 + i];
        }
        out_pts[out_pts_base + tid] = local_pt;
    }
}

torch::Tensor ipa_attention_cuda(
    torch::Tensor q,
    torch::Tensor k,
    torch::Tensor v,
    torch::Tensor q_pts,
    torch::Tensor k_pts,
    torch::Tensor v_pts,
    torch::Tensor pair_bias,
    torch::Tensor head_weights,
    torch::Tensor rotations,
    torch::Tensor translations,
    int N
) {
    int B = q.size(0);
    int h = q.size(2);
    int d = q.size(3);
    int Qp = q_pts.size(3);
    int Vp = v_pts.size(3);
    
    auto options = q.options();
    auto out_scalar = torch::zeros({B, N, h, d}, options);
    auto out_pts = torch::zeros({B, N, h * Vp, 3}, options);
    
    const int block_size = 128;
    dim3 grid(N, B, h);
    // smem: d + Qp*3 + 9 + 3 + N + Vp*3
    size_t smem_size = (d + Qp * 3 + 9 + 3 + N + Vp * 3) * sizeof(float);
    
    ipa_attention_kernel<<<grid, block_size, smem_size>>>(
        q.data_ptr<float>(),
        k.data_ptr<float>(),
        v.data_ptr<float>(),
        q_pts.data_ptr<float>(),
        k_pts.data_ptr<float>(),
        v_pts.data_ptr<float>(),
        pair_bias.data_ptr<float>(),
        head_weights.data_ptr<float>(),
        rotations.data_ptr<float>(),
        translations.data_ptr<float>(),
        out_scalar.data_ptr<float>(),
        out_pts.data_ptr<float>(),
        B, N, h, d, Qp, Vp
    );
    
    C10_CUDA_CHECK(cudaGetLastError());
    
    return torch::cat({out_scalar.reshape({B, N, h * d}), out_pts.reshape({B, N, h * Vp * 3})}, -1);
}
"""

cpp_source = """
torch::Tensor ipa_attention_cuda(
    torch::Tensor q,
    torch::Tensor k,
    torch::Tensor v,
    torch::Tensor q_pts,
    torch::Tensor k_pts,
    torch::Tensor v_pts,
    torch::Tensor pair_bias,
    torch::Tensor head_weights,
    torch::Tensor rotations,
    torch::Tensor translations,
    int N);
"""

ipa_extension = load_inline(
    name='ipa_attention_v3',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['ipa_attention_cuda'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    verbose=True
)


class ModelNew(nn.Module):
    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."""
        global_pts = torch.einsum("bnpc,bncd->bnpd", pts, rotations) + translations.unsqueeze(2)
        return global_pts

    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)

        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)

        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)

        pair_bias_4d = self.pair_bias(pair_repr)  # (B, N, N, h)
        rotations_flat = rotations.reshape(B, N, 9)
        
        result = ipa_extension.ipa_attention_cuda(
            q.contiguous(),
            k.contiguous(),
            v.contiguous(),
            q_pts.contiguous(),
            k_pts.contiguous(),
            v_pts.contiguous(),
            pair_bias_4d.contiguous(),
            self.head_weights.contiguous(),
            rotations_flat.contiguous(),
            translations.contiguous(),
            N
        )
        return self.out_proj(result)