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 qp_w, torch::Tensor kp_w, torch::Tensor vp_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 <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>
#define D 64
#define H 4
#define HD 16
#define NMAX 32
#define PAIRD 32
#define QP 4
#define VP 8
#define OUTIN 160
__device__ __forceinline__ float softplusf_dev(float x){ return x > 20.0f ? x : log1pf(expf(x)); }
__global__ void proj_kernel(const float* __restrict__ node,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__ qp_w,const float* __restrict__ kp_w,const float* __restrict__ vp_w,const float* __restrict__ trans,const float* __restrict__ rot,float* __restrict__ q,float* __restrict__ k,float* __restrict__ v,float* __restrict__ qpts,float* __restrict__ kpts,float* __restrict__ vpts){
    int row=blockIdx.x, tid=threadIdx.x; __shared__ float red[256]; __shared__ float sx[D];
    float sum=0.f; for(int r=tid;r<D;r+=blockDim.x) sum+=node[row*D+r]; red[tid]=sum; __syncthreads();
    for(int s=128;s>0;s>>=1){ if(tid<s) red[tid]+=red[tid+s]; __syncthreads(); } float mean=red[0]*(1.0f/D);
    float ss=0.f; for(int r=tid;r<D;r+=blockDim.x){ float d=node[row*D+r]-mean; ss+=d*d; } red[tid]=ss; __syncthreads();
    for(int s=128;s>0;s>>=1){ if(tid<s) red[tid]+=red[tid+s]; __syncthreads(); } float invstd=rsqrtf(red[0]*(1.0f/D)+1.0e-5f);
    if(tid<D) sx[tid]=(node[row*D+tid]-mean)*invstd*ln_w[tid]+ln_b[tid]; __syncthreads();
    if(tid<192){
        int which=tid/D; int o=tid-which*D; const float* w=(which==0?q_w:(which==1?k_w:v_w)); float acc=0.f;
        for(int r=0;r<D;++r) acc+=sx[r]*w[o*D+r];
        int base=row*H*HD+o; if(which==0) q[base]=acc; else if(which==1) k[base]=acc; else v[base]=acc;
    } else if(tid<256){
        int pidx=tid-192; const float* w; float* out; int point,npoints;
        if(pidx<H*QP){w=qp_w;out=qpts;point=pidx;npoints=H*QP;} else if(pidx<2*H*QP){w=kp_w;out=kpts;point=pidx-H*QP;npoints=H*QP;} else {w=vp_w;out=vpts;point=pidx-2*H*QP;npoints=H*VP;}
        if(point<npoints){
            float loc0=0.f,loc1=0.f,loc2=0.f; int o0=point*3;
            for(int r=0;r<D;++r){ float x=sx[r]; loc0+=x*w[(o0+0)*D+r]; loc1+=x*w[(o0+1)*D+r]; loc2+=x*w[(o0+2)*D+r]; }
            const float* R=rot+row*9; const float* T=trans+row*3; float* dst=out+(row*npoints+point)*3;
            dst[0]=loc0*R[0]+loc1*R[3]+loc2*R[6]+T[0]; dst[1]=loc0*R[1]+loc1*R[4]+loc2*R[7]+T[1]; dst[2]=loc0*R[2]+loc1*R[5]+loc2*R[8]+T[2];
        }
    }
}
__global__ void attn_kernel(const float* __restrict__ q,const float* __restrict__ k,const float* __restrict__ v,const float* __restrict__ qpts,const float* __restrict__ kpts,const float* __restrict__ vpts,const float* __restrict__ pair,const float* __restrict__ pair_w,const float* __restrict__ head_w,float* __restrict__ out_s,float* __restrict__ out_p){
    int tid=threadIdx.x,lane=tid&31,warp=tid>>5,rowidx=blockIdx.x*8+warp; if(rowidx>=2*H*NMAX) return; int i=rowidx%NMAX; int tmp=rowidx/NMAX; int h=tmp%H; int b=tmp/H; int j=lane;
    float dot=0.f; for(int dd=0;dd<HD;++dd) dot+=q[((b*NMAX+i)*H+h)*HD+dd]*k[((b*NMAX+j)*H+h)*HD+dd]; float logit=dot*0.25f;
    float dist=0.f; for(int p=0;p<QP;++p){ int qi=((b*NMAX+i)*(H*QP)+h*QP+p)*3; int kj=((b*NMAX+j)*(H*QP)+h*QP+p)*3; float dx=qpts[qi]-kpts[kj],dy=qpts[qi+1]-kpts[kj+1],dz=qpts[qi+2]-kpts[kj+2]; dist+=dx*dx+dy*dy+dz*dz; }
    logit += -0.5f*softplusf_dev(head_w[h])*dist; float pb=0.f; int pairbase=((b*NMAX+i)*NMAX+j)*PAIRD; for(int p=0;p<PAIRD;++p) pb+=pair[pairbase+p]*pair_w[h*PAIRD+p]; logit+=pb;
    float m=logit; for(int off=16;off>0;off>>=1) m=fmaxf(m,__shfl_down_sync(0xffffffff,m,off)); m=__shfl_sync(0xffffffff,m,0); float e=expf(logit-m); float den=e; for(int off=16;off>0;off>>=1) den+=__shfl_down_sync(0xffffffff,den,off); den=__shfl_sync(0xffffffff,den,0); float a=e/den;
    for(int dd=0;dd<HD;++dd){ float val=a*v[((b*NMAX+j)*H+h)*HD+dd]; for(int off=16;off>0;off>>=1) val+=__shfl_down_sync(0xffffffff,val,off); if(lane==0) out_s[((b*NMAX+i)*H+h)*HD+dd]=val; }
    for(int p=0;p<VP;++p){ for(int c=0;c<3;++c){ float val=a*vpts[((b*NMAX+j)*(H*VP)+h*VP+p)*3+c]; for(int off=16;off>0;off>>=1) val+=__shfl_down_sync(0xffffffff,val,off); if(lane==0) out_p[((b*NMAX+i)*(H*VP)+h*VP+p)*3+c]=val; }}
}
__global__ void out_kernel(const float* __restrict__ out_s,const float* __restrict__ out_p,const float* __restrict__ trans,const float* __restrict__ rot,const float* __restrict__ ow,const float* __restrict__ ob,float* __restrict__ out){
    int row=blockIdx.x,tid=threadIdx.x; const float* R=rot+row*9; const float* T=trans+row*3;
    for(int o=tid;o<D;o+=blockDim.x){ float acc=ob[o]; const float* wrow=ow+o*OUTIN; for(int kf=0;kf<OUTIN;++kf){ float feat; if(kf<H*HD) feat=out_s[row*H*HD+kf]; else { int t=kf-H*HD; int p=t/3; int c=t-p*3; const float* gp=out_p+(row*(H*VP)+p)*3; float cx=gp[0]-T[0],cy=gp[1]-T[1],cz=gp[2]-T[2]; feat=cx*R[c*3]+cy*R[c*3+1]+cz*R[c*3+2]; } acc+=feat*wrow[kf]; } out[row*D+o]=acc; }
}
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 qp_w, torch::Tensor kp_w, torch::Tensor vp_w,torch::Tensor pair_w, torch::Tensor head_w,torch::Tensor out_w, torch::Tensor out_b){
    TORCH_CHECK(node.is_cuda(),"cuda only"); int B=node.size(0),N=node.size(1); TORCH_CHECK(B==2&&N==32&&node.size(2)==64,"specialized"); auto opts=node.options(); auto q=torch::empty({B,N,H,HD},opts),k=torch::empty({B,N,H,HD},opts),v=torch::empty({B,N,H,HD},opts); auto qpts=torch::empty({B,N,H*QP,3},opts),kpts=torch::empty({B,N,H*QP,3},opts),vpts=torch::empty({B,N,H*VP,3},opts); auto outs=torch::empty({B,N,H,HD},opts); auto outp=torch::empty({B,N,H*VP,3},opts); auto out=torch::empty({B,N,D},opts); cudaStream_t stream=at::cuda::getCurrentCUDAStream();
    proj_kernel<<<B*N,256,0,stream>>>(node.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>(),qp_w.data_ptr<float>(),kp_w.data_ptr<float>(),vp_w.data_ptr<float>(),trans.data_ptr<float>(),rot.data_ptr<float>(),q.data_ptr<float>(),k.data_ptr<float>(),v.data_ptr<float>(),qpts.data_ptr<float>(),kpts.data_ptr<float>(),vpts.data_ptr<float>());
    attn_kernel<<<(B*H*N+7)/8,256,0,stream>>>(q.data_ptr<float>(),k.data_ptr<float>(),v.data_ptr<float>(),qpts.data_ptr<float>(),kpts.data_ptr<float>(),vpts.data_ptr<float>(),pair.data_ptr<float>(),pair_w.data_ptr<float>(),head_w.data_ptr<float>(),outs.data_ptr<float>(),outp.data_ptr<float>());
    out_kernel<<<B*N,128,0,stream>>>(outs.data_ptr<float>(),outp.data_ptr<float>(),trans.data_ptr<float>(),rot.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_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)); self.out_proj=nn.Linear(num_heads*d+num_heads*num_value_points*3,node_dim); self.layer_norm=nn.LayerNorm(node_dim)
    def forward(self,node_repr,pair_repr,translations,rotations):
        return _ext.ipa_forward(node_repr,pair_repr,translations,rotations,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)