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


class Model(nn.Module):
    """
    Column-wise gated self-attention over MSA representations.
    For each residue position, attention is computed across the MSA depth
    (sequence) dimension.  This lets the model share information between
    different aligned sequences at the same position.
    """

    def __init__(self, msa_dim, num_heads, dropout=0.0):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = msa_dim // num_heads
        assert msa_dim % num_heads == 0

        self.layer_norm = nn.LayerNorm(msa_dim)
        self.query = nn.Linear(msa_dim, msa_dim, bias=False)
        self.key = nn.Linear(msa_dim, msa_dim, bias=False)
        self.value = nn.Linear(msa_dim, msa_dim, bias=False)
        self.gate = nn.Linear(msa_dim, msa_dim)
        self.out_proj = nn.Linear(msa_dim, msa_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, msa_repr: torch.Tensor, msa_mask: torch.Tensor) -> torch.Tensor:
        """
        Args:
            msa_repr: (B, S, N, msa_dim)
            msa_mask: (B, S, N)
        Returns:
            updated MSA (B, S, N, msa_dim)
        """
        B, S, N, D = msa_repr.shape
        h = self.num_heads
        d = self.head_dim

        # Transpose so column (residue position) is the batch dim for attention
        x = msa_repr.permute(0, 2, 1, 3)  # (B, N, S, D)
        x = self.layer_norm(x)

        q = self.query(x).view(B, N, S, h, d).permute(0, 1, 3, 2, 4)
        k = self.key(x).view(B, N, S, h, d).permute(0, 1, 3, 2, 4)
        v = self.value(x).view(B, N, S, h, d).permute(0, 1, 3, 2, 4)
        g = torch.sigmoid(self.gate(x)).view(B, N, S, h, d).permute(0, 1, 3, 2, 4)

        attn = torch.einsum("bnhid,bnhjd->bnhij", q, k) / math.sqrt(d)

        col_mask = msa_mask.permute(0, 2, 1)  # (B, N, S)
        mask = col_mask.view(B, N, 1, 1, S).expand_as(attn)
        attn = attn.masked_fill(mask == 0, float("-inf"))
        attn = F.softmax(attn, dim=-1)
        attn = self.dropout(attn)

        out = torch.einsum("bnhij,bnhjd->bnhid", attn, v)
        out = g * out
        out = out.permute(0, 1, 3, 2, 4).reshape(B, N, S, D)
        out = self.out_proj(out)
        return out.permute(0, 2, 1, 3)  # back to (B, S, N, D)


msa_dim = 64
num_heads = 4
seq_len = 32
msa_depth = 8
batch_size = 2


def get_inputs():
    p = popcorn_pri
    msa_repr = torch.randn(p.jitter_int(batch_size), p.jitter_int(msa_depth), p.jitter_int(seq_len), msa_dim)
    msa_mask = torch.ones(p.jitter_int(batch_size), p.jitter_int(msa_depth), p.jitter_int(seq_len))
    return [msa_repr, msa_mask]


def get_init_inputs():
    return [msa_dim, num_heads]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_src = r'''
#include <torch/types.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>
#define SS 8
#define DD 64
#define HH 4
#define HD 16
#define NW 32
__inline__ __device__ float warp_sum(float v){ unsigned m=0xffffffffu; v+=__shfl_down_sync(m,v,16); v+=__shfl_down_sync(m,v,8); v+=__shfl_down_sync(m,v,4); v+=__shfl_down_sync(m,v,2); v+=__shfl_down_sync(m,v,1); return v; }
__global__ __launch_bounds__(1024,1) void col_attn_kernel(const float* __restrict__ inp,const float* __restrict__ mask,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__ gate_w,const float* __restrict__ gate_b,const float* __restrict__ out_w,const float* __restrict__ out_b,float* __restrict__ out,int B,int S,int N,int D){
 int col=blockIdx.x,b=col/N,n=col-b*N,tid=threadIdx.x,lane=tid&31,warp=tid>>5;
 __shared__ float xn[SS*DD],q[SS*DD],k[SS*DD],v[SS*DD],g[SS*DD],att[HH*SS*SS],ao[SS*DD];
 if(warp<SS){int s=warp; float sum=0.f; for(int d=lane;d<DD;d+=32) sum+=inp[((b*S+s)*N+n)*D+d]; sum=warp_sum(sum); float mean=__shfl_sync(0xffffffffu,sum*(1.f/DD),0); float vs=0.f; for(int d=lane;d<DD;d+=32){float z=inp[((b*S+s)*N+n)*D+d]-mean; vs+=z*z;} vs=warp_sum(vs); float rstd=__shfl_sync(0xffffffffu,rsqrtf(vs*(1.f/DD)+1.e-5f),0); for(int d=lane;d<DD;d+=32){float z=(inp[((b*S+s)*N+n)*D+d]-mean)*rstd; xn[s*DD+d]=z*ln_w[d]+ln_b[d];}}
 __syncthreads();
 for(int idx=warp; idx<SS*DD; idx+=NW){int s=idx/DD,od=idx-s*DD; float qs=0.f,ks=0.f,vs=0.f,gs=0.f; for(int id=lane;id<DD;id+=32){float x=xn[s*DD+id]; qs+=x*q_w[od*DD+id]; ks+=x*k_w[od*DD+id]; vs+=x*v_w[od*DD+id]; gs+=x*gate_w[od*DD+id];} qs=warp_sum(qs); ks=warp_sum(ks); vs=warp_sum(vs); gs=warp_sum(gs); if(lane==0){q[idx]=qs; k[idx]=ks; v[idx]=vs; float z=gs+gate_b[od]; g[idx]=1.f/(1.f+expf(-z));}}
 __syncthreads();
 for(int hi=warp; hi<HH*SS; hi+=NW){int h=hi/SS,i=hi-h*SS; float scores[SS]; float mx=-3.402823466e38f; for(int j=0;j<SS;++j){float dot=0.f; if(lane<HD){int od=h*HD+lane; dot=q[i*DD+od]*k[j*DD+od];} dot=warp_sum(dot); if(lane==0){float m=mask[(b*S+j)*N+n]; float sc=(m==0.f)?-3.402823466e38f:dot*0.25f; scores[j]=sc; mx=fmaxf(mx,sc);}} if(lane==0){float den=0.f; for(int j=0;j<SS;++j){float e=expf(scores[j]-mx); scores[j]=e; den+=e;} float inv=1.f/den; for(int j=0;j<SS;++j) att[(h*SS+i)*SS+j]=scores[j]*inv;}}
 __syncthreads();
 if(tid<SS*DD){int idx=tid,i=idx/DD,od=idx-i*DD,h=od/HD; float sum=0.f; for(int j=0;j<SS;++j) sum+=att[(h*SS+i)*SS+j]*v[j*DD+od]; ao[idx]=sum*g[idx];}
 __syncthreads();
 for(int idx=warp; idx<SS*DD; idx+=NW){int s=idx/DD,od=idx-s*DD; float sum=0.f; for(int id=lane;id<DD;id+=32) sum+=ao[s*DD+id]*out_w[od*DD+id]; sum=warp_sum(sum); if(lane==0) out[((b*S+s)*N+n)*D+od]=sum+out_b[od];}
}
torch::Tensor col_attn_forward(torch::Tensor inp,torch::Tensor mask,torch::Tensor ln_w,torch::Tensor ln_b,torch::Tensor q_w,torch::Tensor k_w,torch::Tensor v_w,torch::Tensor gate_w,torch::Tensor gate_b,torch::Tensor out_w,torch::Tensor out_b){int B=inp.size(0),S=inp.size(1),N=inp.size(2),D=inp.size(3); auto output=torch::empty_like(inp); col_attn_kernel<<<B*N,1024,0,at::cuda::getCurrentCUDAStream()>>>(inp.data_ptr<float>(),mask.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>(),gate_w.data_ptr<float>(),gate_b.data_ptr<float>(),out_w.data_ptr<float>(),out_b.data_ptr<float>(),output.data_ptr<float>(),B,S,N,D); return output;}
'''
cpp_src = '#include <torch/extension.h>\ntorch::Tensor col_attn_forward(torch::Tensor inp, torch::Tensor mask, torch::Tensor ln_w, torch::Tensor ln_b, torch::Tensor q_w, torch::Tensor k_w, torch::Tensor v_w, torch::Tensor gate_w, torch::Tensor gate_b, torch::Tensor out_w, torch::Tensor out_b);\n'
_ext=load_inline(name='col_attn_ext_submit_final',cpp_sources=cpp_src,cuda_sources=cuda_src,functions=['col_attn_forward'],extra_cuda_cflags=['-O3','--use_fast_math'],extra_cflags=['-O3'],verbose=False)
class ModelNew(nn.Module):
 def __init__(self, msa_dim, num_heads, dropout=0.0):
  super().__init__(); self.num_heads=num_heads; self.head_dim=msa_dim//num_heads; assert msa_dim%num_heads==0
  self.layer_norm=nn.LayerNorm(msa_dim); self.query=nn.Linear(msa_dim,msa_dim,bias=False); self.key=nn.Linear(msa_dim,msa_dim,bias=False); self.value=nn.Linear(msa_dim,msa_dim,bias=False); self.gate=nn.Linear(msa_dim,msa_dim); self.out_proj=nn.Linear(msa_dim,msa_dim); self.dropout=nn.Dropout(dropout)
 def forward(self, msa_repr, msa_mask):
  return _ext.col_attn_forward(msa_repr.contiguous(), msa_mask.contiguous(), self.layer_norm.weight, self.layer_norm.bias, self.query.weight, self.key.weight, self.value.weight, self.gate.weight, self.gate.bias, self.out_proj.weight, self.out_proj.bias)