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
from torch.utils.cpp_extension import load_inline

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor ipa_forward(torch::Tensor node, torch::Tensor pair, torch::Tensor trans, torch::Tensor rot,
                          torch::Tensor ln_w, torch::Tensor ln_b,
                          torch::Tensor q_w, torch::Tensor k_w, torch::Tensor v_w,
                          torch::Tensor qpts_w, torch::Tensor kpts_w, torch::Tensor vpts_w,
                          torch::Tensor pair_w, torch::Tensor head_w,
                          torch::Tensor out_w, torch::Tensor out_b);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>
#define D 64
#define H 4
#define HD 16
#define NSEQ 32
#define PD 32
#define QP 4
#define VP 8
#define CTX 160
#define QOFF 0
#define KOFF (H*HD)
#define VOFF (2*H*HD)
#define QPOFF (3*H*HD)
#define KPOFF (QPOFF + H*QP*3)
#define VPOFF (KPOFF + H*QP*3)
#define STRIDE (VPOFF + H*VP*3)

__global__ void proj_kernel(const float* __restrict__ node,
                            const float* __restrict__ trans,
                            const float* __restrict__ rot,
                            const float* __restrict__ ln_w,
                            const float* __restrict__ ln_b,
                            const float* __restrict__ q_w,
                            const float* __restrict__ k_w,
                            const float* __restrict__ v_w,
                            const float* __restrict__ qpts_w,
                            const float* __restrict__ kpts_w,
                            const float* __restrict__ vpts_w,
                            float* __restrict__ buf) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    __shared__ float sh[128];
    __shared__ float x[D];
    float a = (tid < D) ? node[row*D + tid] : 0.f;
    sh[tid] = a;
    __syncthreads();
    for (int s=64; s>0; s>>=1) { if (tid < s) sh[tid] += sh[tid+s]; __syncthreads(); }
    float mean = sh[0] * (1.0f/D);
    float diff = (tid < D) ? (node[row*D + tid] - mean) : 0.f;
    sh[tid] = (tid < D) ? diff*diff : 0.f;
    __syncthreads();
    for (int s=64; s>0; s>>=1) { if (tid < s) sh[tid] += sh[tid+s]; __syncthreads(); }
    float invstd = rsqrtf(sh[0] * (1.0f/D) + 1.0e-5f);
    if (tid < D) x[tid] = diff * invstd * ln_w[tid] + ln_b[tid];
    __syncthreads();
    float* out = buf + row*STRIDE;
    const float* R = rot + row*9;
    const float* T = trans + row*3;
    for (int oc=tid; oc<3*H*HD; oc+=blockDim.x) {
        int which = oc >> 6;
        int o = oc & 63;
        const float* W = (which==0 ? q_w : (which==1 ? k_w : v_w));
        float s=0.f;
        #pragma unroll
        for (int c=0; c<D; ++c) s += x[c] * W[o*D+c];
        out[which*H*HD + o] = s;
    }
    for (int pidx=tid; pidx<H*QP; pidx+=blockDim.x) {
        float l0=0.f,l1=0.f,l2=0.f;
        const float* W = qpts_w + pidx*3*D;
        #pragma unroll
        for (int c=0; c<D; ++c) { float xv=x[c]; l0 += xv*W[c]; l1 += xv*W[D+c]; l2 += xv*W[2*D+c]; }
        out[QPOFF + pidx*3 + 0] = l0*R[0] + l1*R[3] + l2*R[6] + T[0];
        out[QPOFF + pidx*3 + 1] = l0*R[1] + l1*R[4] + l2*R[7] + T[1];
        out[QPOFF + pidx*3 + 2] = l0*R[2] + l1*R[5] + l2*R[8] + T[2];
        l0=0.f; l1=0.f; l2=0.f; W = kpts_w + pidx*3*D;
        #pragma unroll
        for (int c=0; c<D; ++c) { float xv=x[c]; l0 += xv*W[c]; l1 += xv*W[D+c]; l2 += xv*W[2*D+c]; }
        out[KPOFF + pidx*3 + 0] = l0*R[0] + l1*R[3] + l2*R[6] + T[0];
        out[KPOFF + pidx*3 + 1] = l0*R[1] + l1*R[4] + l2*R[7] + T[1];
        out[KPOFF + pidx*3 + 2] = l0*R[2] + l1*R[5] + l2*R[8] + T[2];
    }
    for (int pidx=tid; pidx<H*VP; pidx+=blockDim.x) {
        float l0=0.f,l1=0.f,l2=0.f;
        const float* W = vpts_w + pidx*3*D;
        #pragma unroll
        for (int c=0; c<D; ++c) { float xv=x[c]; l0 += xv*W[c]; l1 += xv*W[D+c]; l2 += xv*W[2*D+c]; }
        out[VPOFF + pidx*3 + 0] = l0*R[0] + l1*R[3] + l2*R[6] + T[0];
        out[VPOFF + pidx*3 + 1] = l0*R[1] + l1*R[4] + l2*R[7] + T[1];
        out[VPOFF + pidx*3 + 2] = l0*R[2] + l1*R[5] + l2*R[8] + T[2];
    }
}

__device__ __forceinline__ float softplus_dev(float x) { return x > 20.0f ? x : log1pf(expf(x)); }

__global__ void attn_out_kernel(const float* __restrict__ buf,
                                const float* __restrict__ pair,
                                const float* __restrict__ trans,
                                const float* __restrict__ rot,
                                const float* __restrict__ pair_w,
                                const float* __restrict__ head_w,
                                const float* __restrict__ out_w,
                                const float* __restrict__ out_b,
                                float* __restrict__ out) {
    int row_i = blockIdx.x;
    int b = row_i / NSEQ;
    int i = row_i - b*NSEQ;
    int tid = threadIdx.x;
    __shared__ float logits[H*NSEQ];
    __shared__ float red[128];
    __shared__ float ctx[CTX];
    const float* rowbuf = buf + row_i*STRIDE;
    for (int h=0; h<H; ++h) {
        if (tid < NSEQ) {
            int j = tid;
            const float* jb = buf + (b*NSEQ + j)*STRIDE;
            const float* q = rowbuf + QOFF + h*HD;
            const float* k = jb + KOFF + h*HD;
            float s=0.f;
            #pragma unroll
            for (int d=0; d<HD; ++d) s += q[d]*k[d];
            s *= 0.25f;
            const float* qp = rowbuf + QPOFF + h*QP*3;
            const float* kp = jb + KPOFF + h*QP*3;
            float dist=0.f;
            #pragma unroll
            for (int p=0; p<QP*3; ++p) { float df=qp[p]-kp[p]; dist += df*df; }
            s += -0.5f * softplus_dev(head_w[h]) * dist;
            const float* pr = pair + ((b*NSEQ+i)*NSEQ+j)*PD;
            const float* pw = pair_w + h*PD;
            float pb=0.f;
            #pragma unroll
            for (int d=0; d<PD; ++d) pb += pr[d]*pw[d];
            s += pb;
            logits[h*NSEQ + j] = s;
            red[tid] = s;
        } else {
            red[tid] = -3.402823466e38f;
        }
        __syncthreads();
        for (int stride=64; stride>0; stride>>=1) { if (tid < stride) red[tid] = fmaxf(red[tid], red[tid+stride]); __syncthreads(); }
        float m = red[0];
        float e = 0.f;
        if (tid < NSEQ) { e = expf(logits[h*NSEQ+tid] - m); logits[h*NSEQ+tid] = e; }
        red[tid] = (tid < NSEQ) ? e : 0.f;
        __syncthreads();
        for (int stride=64; stride>0; stride>>=1) { if (tid < stride) red[tid] += red[tid+stride]; __syncthreads(); }
        float invsum = 1.0f / red[0];
        if (tid < NSEQ) logits[h*NSEQ+tid] *= invsum;
        __syncthreads();
        for (int f=tid; f<HD; f+=blockDim.x) {
            float acc=0.f;
            #pragma unroll
            for (int j=0; j<NSEQ; ++j) acc += logits[h*NSEQ+j] * buf[(b*NSEQ+j)*STRIDE + VOFF + h*HD + f];
            ctx[h*HD + f] = acc;
        }
        for (int f=tid; f<VP*3; f+=blockDim.x) {
            int p = f/3;
            int c = f - p*3;
            const float* R = rot + row_i*9;
            const float* T = trans + row_i*3;
            float acc=0.f;
            #pragma unroll
            for (int j=0; j<NSEQ; ++j) {
                const float* gp = buf + (b*NSEQ+j)*STRIDE + VPOFF + (h*VP+p)*3;
                float dot = gp[0]*R[c*3+0] + gp[1]*R[c*3+1] + gp[2]*R[c*3+2];
                acc += logits[h*NSEQ+j] * dot;
            }
            float tdot = T[0]*R[c*3+0] + T[1]*R[c*3+1] + T[2]*R[c*3+2];
            ctx[H*HD + h*VP*3 + f] = acc - tdot;
        }
        __syncthreads();
    }
    for (int o=tid; o<D; o+=blockDim.x) {
        float s = out_b[o];
        const float* W = out_w + o*CTX;
        #pragma unroll
        for (int f=0; f<CTX; ++f) s += ctx[f]*W[f];
        out[row_i*D + o] = s;
    }
}

torch::Tensor ipa_forward(torch::Tensor node, torch::Tensor pair, torch::Tensor trans, torch::Tensor rot,
                          torch::Tensor ln_w, torch::Tensor ln_b,
                          torch::Tensor q_w, torch::Tensor k_w, torch::Tensor v_w,
                          torch::Tensor qpts_w, torch::Tensor kpts_w, torch::Tensor vpts_w,
                          torch::Tensor pair_w, torch::Tensor head_w,
                          torch::Tensor out_w, torch::Tensor out_b) {
    int B = node.size(0);
    int N = node.size(1);
    TORCH_CHECK(N == NSEQ && node.size(2) == D, "fixed shape only");
    int rows = B * N;
    auto opts = node.options();
    auto buf = torch::empty({rows, STRIDE}, opts);
    auto out = torch::empty({B,N,D}, opts);
    proj_kernel<<<rows, 128>>>(node.data_ptr<float>(), trans.data_ptr<float>(), rot.data_ptr<float>(),
        ln_w.data_ptr<float>(), ln_b.data_ptr<float>(), q_w.data_ptr<float>(), k_w.data_ptr<float>(), v_w.data_ptr<float>(),
        qpts_w.data_ptr<float>(), kpts_w.data_ptr<float>(), vpts_w.data_ptr<float>(), buf.data_ptr<float>());
    attn_out_kernel<<<rows, 128>>>(buf.data_ptr<float>(), pair.data_ptr<float>(), trans.data_ptr<float>(), rot.data_ptr<float>(),
        pair_w.data_ptr<float>(), head_w.data_ptr<float>(), out_w.data_ptr<float>(), out_b.data_ptr<float>(), out.data_ptr<float>());
    return out;
}
'''

_ext = load_inline(name='ipa_ext_v2_final', cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=['ipa_forward'], extra_cuda_cflags=['-O3'], verbose=False)

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 forward(self, node_repr, pair_repr, translations, rotations):
        return _ext.ipa_forward(node_repr.contiguous(), pair_repr.contiguous(), translations.contiguous(), rotations.contiguous(),
                                self.layer_norm.weight, self.layer_norm.bias,
                                self.q_proj.weight, self.k_proj.weight, self.v_proj.weight,
                                self.q_pts.weight, self.k_pts.weight, self.v_pts.weight,
                                self.pair_bias.weight, self.head_weights,
                                self.out_proj.weight, self.out_proj.bias)