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


class Model(nn.Module):
    """
    Row-wise gated self-attention over Multiple Sequence Alignment (MSA)
    representations.  Each row of the MSA is an independent sequence; for
    each row, standard multi-head attention is performed across residue
    positions, with an additive bias from the pair representation.
    """

    def __init__(self, msa_dim, pair_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.pair_bias = nn.Linear(pair_dim, num_heads, bias=False)
        self.out_proj = nn.Linear(msa_dim, msa_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(
        self,
        msa_repr: torch.Tensor,
        pair_repr: torch.Tensor,
        msa_mask: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            msa_repr:  (B, S, N, msa_dim)
            pair_repr: (B, N, N, pair_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

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

        # Pair bias: shared across MSA rows
        bias = self.pair_bias(pair_repr).permute(0, 3, 1, 2)  # (B, h, N, N)

        attn = torch.einsum("bshid,bshjd->bshij", q, k) / math.sqrt(d)
        attn = attn + bias.unsqueeze(1)

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

        out = torch.einsum("bshij,bshjd->bshid", attn, v)
        out = g * out
        out = out.permute(0, 1, 3, 2, 4).reshape(B, S, N, D)
        return self.out_proj(out)


msa_dim = 64
pair_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)
    pair_repr = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(seq_len), pair_dim)
    msa_mask = torch.ones(p.jitter_int(batch_size), p.jitter_int(msa_depth), p.jitter_int(seq_len))
    return [msa_repr, pair_repr, msa_mask]


def get_init_inputs():
    return [msa_dim, pair_dim, num_heads]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp = r'''
#include <torch/extension.h>
torch::Tensor row_attn_forward(torch::Tensor msa, torch::Tensor pair, torch::Tensor mask,
    torch::Tensor ln_w, torch::Tensor ln_b,
    torch::Tensor q_wt, torch::Tensor k_wt, torch::Tensor v_wt,
    torch::Tensor gate_wt, torch::Tensor gate_b,
    torch::Tensor pair_wt,
    torch::Tensor out_wt, torch::Tensor out_b,
    torch::Tensor work, torch::Tensor out);
'''

_cuda = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
#define D 64
#define H 4
#define HD 16
#define N 32
#define S 8
#define BATCH 2
#define TOKENS (BATCH*S*N)
#define QSIZE (BATCH*S*H*N*HD)
#define BSIZE (BATCH*H*N*N)
__device__ __forceinline__ float sigmoidf_fast(float x){return 1.0f/(1.0f+expf(-x));}

__global__ void ln_proj_kernel(const float* __restrict__ msa,const float* __restrict__ ln_w,const float* __restrict__ ln_b,const float* __restrict__ q_wt,const float* __restrict__ k_wt,const float* __restrict__ v_wt,const float* __restrict__ gate_wt,const float* __restrict__ gate_b,float* __restrict__ q,float* __restrict__ k,float* __restrict__ v,float* __restrict__ g){
    int token=blockIdx.x, tid=threadIdx.x; __shared__ float sh[66]; float x0=msa[token*D+tid]; sh[tid]=x0;
    __syncthreads(); if(tid<32)sh[tid]+=sh[tid+32];
    __syncthreads(); if(tid<16)sh[tid]+=sh[tid+16];
    __syncthreads(); if(tid<8)sh[tid]+=sh[tid+8];
    __syncthreads(); if(tid<4)sh[tid]+=sh[tid+4];
    __syncthreads(); if(tid<2)sh[tid]+=sh[tid+2];
    __syncthreads(); if(tid==0){sh[0]+=sh[1]; sh[64]=sh[0]*0.015625f;} __syncthreads();
    float mean=sh[64]; float a=x0-mean; sh[tid]=a*a;
    __syncthreads(); if(tid<32)sh[tid]+=sh[tid+32];
    __syncthreads(); if(tid<16)sh[tid]+=sh[tid+16];
    __syncthreads(); if(tid<8)sh[tid]+=sh[tid+8];
    __syncthreads(); if(tid<4)sh[tid]+=sh[tid+4];
    __syncthreads(); if(tid<2)sh[tid]+=sh[tid+2];
    __syncthreads(); if(tid==0){sh[0]+=sh[1]; sh[65]=rsqrtf(sh[0]*0.015625f+1.0e-5f);} __syncthreads();
    sh[tid]=a*sh[65]*ln_w[tid]+ln_b[tid]; __syncthreads();
    int out=tid; float qs=0.f,ks=0.f,vs=0.f,gs=gate_b[out];
    #pragma unroll
    for(int i=0;i<D;i++){float x=sh[i]; qs+=x*q_wt[i*D+out]; ks+=x*k_wt[i*D+out]; vs+=x*v_wt[i*D+out]; gs+=x*gate_wt[i*D+out];}
    int n=token&31, bs=token>>5, head=out>>4, dd=out&15; int idx=(((bs*H+head)*N+n)*HD+dd); q[idx]=qs; k[idx]=ks; v[idx]=vs; g[idx]=sigmoidf_fast(gs);
}

__global__ void pair_bias_kernel(const float* __restrict__ pair,const float* __restrict__ pair_wt,float* __restrict__ bias){
    int idx=blockIdx.x*blockDim.x+threadIdx.x; if(idx>=BSIZE)return;
    int t=idx; int j=t&31; t>>=5; int i=t&31; t>>=5; int h=t&3; t>>=2; int b=t;
    const float* p=pair+((b*N+i)*N+j)*D; float s=0.f;
    #pragma unroll
    for(int c=0;c<D;c++) s+=p[c]*pair_wt[c*H+h];
    bias[idx]=s;
}

__global__ void attn_out_kernel(const float* __restrict__ q,const float* __restrict__ k,const float* __restrict__ v,const float* __restrict__ g,const float* __restrict__ bias,const float* __restrict__ mask,const float* __restrict__ out_wt,const float* __restrict__ out_b,float* __restrict__ out){
    int token=blockIdx.x, tid=threadIdx.x; int i=token&31, bs=token>>5, b=bs>>3; __shared__ float scores[H*N]; __shared__ float temp[D];
    if(tid<H*N){
        int h=tid>>5, j=tid&31; const float* qp=q+(((bs*H+h)*N+i)*HD); const float* kp=k+(((bs*H+h)*N+j)*HD); float dot=0.f;
        #pragma unroll
        for(int dd=0;dd<HD;dd++) dot+=qp[dd]*kp[dd];
        float sc=dot*0.25f+bias[((b*H+h)*N+i)*N+j]; float m=mask[bs*N+j]; scores[tid]=(m==0.0f)?-INFINITY:sc;
    }
    __syncthreads();
    if(tid<H){
        int base=tid*N; float mx=scores[base];
        #pragma unroll
        for(int j=1;j<N;j++) mx=fmaxf(mx,scores[base+j]);
        float sum=0.f;
        #pragma unroll
        for(int j=0;j<N;j++){float e=expf(scores[base+j]-mx); scores[base+j]=e; sum+=e;}
        float inv=1.0f/sum;
        #pragma unroll
        for(int j=0;j<N;j++) scores[base+j]*=inv;
    }
    __syncthreads();
    if(tid<D){
        int h=tid>>4, dd=tid&15; float acc=0.f;
        #pragma unroll
        for(int j=0;j<N;j++) acc+=scores[h*N+j]*v[(((bs*H+h)*N+j)*HD+dd)];
        int src=(((bs*H+h)*N+i)*HD+dd); temp[tid]=acc*g[src];
    }
    __syncthreads();
    if(tid<D){
        int o=tid; float acc=out_b[o];
        #pragma unroll
        for(int c=0;c<D;c++) acc+=temp[c]*out_wt[c*D+o];
        out[token*D+o]=acc;
    }
}

torch::Tensor row_attn_forward(torch::Tensor msa, torch::Tensor pair, torch::Tensor mask,torch::Tensor ln_w, torch::Tensor ln_b,torch::Tensor q_wt, torch::Tensor k_wt, torch::Tensor v_wt,torch::Tensor gate_wt, torch::Tensor gate_b,torch::Tensor pair_wt,torch::Tensor out_wt, torch::Tensor out_b,torch::Tensor work, torch::Tensor out){
    float* base=work.data_ptr<float>(); float* q=base; float* k=q+QSIZE; float* v=k+QSIZE; float* g=v+QSIZE; float* bias=g+QSIZE;
    ln_proj_kernel<<<TOKENS,64>>>(msa.data_ptr<float>(),ln_w.data_ptr<float>(),ln_b.data_ptr<float>(),q_wt.data_ptr<float>(),k_wt.data_ptr<float>(),v_wt.data_ptr<float>(),gate_wt.data_ptr<float>(),gate_b.data_ptr<float>(),q,k,v,g);
    pair_bias_kernel<<<(BSIZE+255)/256,256>>>(pair.data_ptr<float>(),pair_wt.data_ptr<float>(),bias);
    attn_out_kernel<<<TOKENS,128>>>(q,k,v,g,bias,mask.data_ptr<float>(),out_wt.data_ptr<float>(),out_b.data_ptr<float>(),out.data_ptr<float>());
    return out;
}
'''

_ext=load_inline(name='row_attn_ext_final', cpp_sources=_cpp, cuda_sources=_cuda, functions=['row_attn_forward'], extra_cuda_cflags=['-O3','--use_fast_math'], extra_cflags=['-O3'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, msa_dim, pair_dim, num_heads, dropout=0.0):
        super().__init__()
        self.num_heads=num_heads; self.head_dim=msa_dim//num_heads
        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.pair_bias=nn.Linear(pair_dim,num_heads,bias=False)
        self.out_proj=nn.Linear(msa_dim,msa_dim)
        self.dropout=nn.Dropout(dropout)
        self.register_buffer('_work', torch.empty(0), persistent=False)
        self.register_buffer('_out', torch.empty(0), persistent=False)
        self.register_buffer('_qwt', torch.empty(0), persistent=False)
        self.register_buffer('_kwt', torch.empty(0), persistent=False)
        self.register_buffer('_vwt', torch.empty(0), persistent=False)
        self.register_buffer('_gwt', torch.empty(0), persistent=False)
        self.register_buffer('_pwt', torch.empty(0), persistent=False)
        self.register_buffer('_owt', torch.empty(0), persistent=False)
    def _prep(self, x):
        if self._work.numel()!=139264 or self._work.device!=x.device:
            self._work=torch.empty((139264,),device=x.device,dtype=x.dtype)
        if self._out.numel()!=x.numel() or self._out.device!=x.device:
            self._out=torch.empty_like(x)
        if self._qwt.numel()!=4096 or self._qwt.device!=x.device:
            self._qwt=self.query.weight.detach().t().contiguous()
            self._kwt=self.key.weight.detach().t().contiguous()
            self._vwt=self.value.weight.detach().t().contiguous()
            self._gwt=self.gate.weight.detach().t().contiguous()
            self._pwt=self.pair_bias.weight.detach().t().contiguous()
            self._owt=self.out_proj.weight.detach().t().contiguous()
    def forward(self, msa_repr, pair_repr, msa_mask):
        self._prep(msa_repr)
        return _ext.row_attn_forward(msa_repr,pair_repr,msa_mask,self.layer_norm.weight,self.layer_norm.bias,self._qwt,self._kwt,self._vwt,self._gwt,self.gate.bias,self._pwt,self._owt,self.out_proj.bias,self._work,self._out)