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

cpp_src = 'torch::Tensor col_attn_forward(torch::Tensor msa_repr, torch::Tensor msa_mask, torch::Tensor ln_w, torch::Tensor ln_b, torch::Tensor qw, torch::Tensor kw, torch::Tensor vw, torch::Tensor gw, torch::Tensor gb, torch::Tensor ow, torch::Tensor ob);'

cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
#define D 64
#define S_MAX 8
#define HD 16
__device__ __forceinline__ float sigmoidf_fast(float x){ return 1.0f / (1.0f + expf(-x)); }
__global__ void col_attn_kernel(const float* __restrict__ x_in, const float* __restrict__ mask, const float* __restrict__ ln_w, const float* __restrict__ ln_b, const float* __restrict__ qw, const float* __restrict__ kw, const float* __restrict__ vw, const float* __restrict__ gw, const float* __restrict__ gb, const float* __restrict__ ow, const float* __restrict__ ob, float* __restrict__ y, int B, int S, int N){
    int col=blockIdx.x, b=col/N, n=col-b*N, tid=threadIdx.x;
    __shared__ float xn[S_MAX*D], q[S_MAX*D], k[S_MAX*D], v[S_MAX*D], g[S_MAX*D], tmp[S_MAX*D];
    if(tid < S){
        int s=tid; const float* xp_in=&x_in[((b*S+s)*N+n)*D]; float sum=0.f;
        #pragma unroll
        for(int i=0;i<D;++i) sum += xp_in[i];
        float mean=sum*(1.0f/64.0f); float ss=0.f;
        #pragma unroll
        for(int i=0;i<D;++i){ float d=xp_in[i]-mean; ss += d*d; }
        float r=rsqrtf(ss*(1.0f/64.0f)+1.0e-5f);
        #pragma unroll
        for(int i=0;i<D;++i) xn[s*D+i]=(xp_in[i]-mean)*r*ln_w[i]+ln_b[i];
    }
    __syncthreads();
    for(int idx=tid; idx<S*D; idx+=blockDim.x){
        int s=idx>>6, od=idx&63; float qs=0.f,ks=0.f,vs=0.f,gs=gb[od]; const float* xp=&xn[s*D]; const float* qwp=&qw[od*D]; const float* kwp=&kw[od*D]; const float* vwp=&vw[od*D]; const float* gwp=&gw[od*D];
        #pragma unroll
        for(int in=0; in<D; ++in){ float xv=xp[in]; qs=fmaf(xv,qwp[in],qs); ks=fmaf(xv,kwp[in],ks); vs=fmaf(xv,vwp[in],vs); gs=fmaf(xv,gwp[in],gs); }
        q[idx]=qs; k[idx]=ks; v[idx]=vs; g[idx]=sigmoidf_fast(gs);
    }
    __syncthreads();
    const float scale=0.25f;
    for(int idx=tid; idx<S*D; idx+=blockDim.x){
        int s=idx>>6, dglob=idx&63, h=dglob>>4; float scores[S_MAX]; float m=-3.402823466e38f;
        #pragma unroll
        for(int j=0;j<S_MAX;++j){
            float sc=-3.402823466e38f;
            if(j<S && mask[(b*S+j)*N+n] != 0.0f){
                float dot=0.f; int bq=s*D+h*HD, bk=j*D+h*HD;
                #pragma unroll
                for(int kk=0;kk<HD;++kk) dot=fmaf(q[bq+kk],k[bk+kk],dot);
                sc=dot*scale;
            }
            scores[j]=sc; if(sc>m)m=sc;
        }
        float denom=0.f;
        #pragma unroll
        for(int j=0;j<S_MAX;++j){ float e=(j<S)?expf(scores[j]-m):0.f; scores[j]=e; denom+=e; }
        float invd=1.0f/denom, acc=0.f;
        #pragma unroll
        for(int j=0;j<S_MAX;++j) if(j<S) acc=fmaf(scores[j]*invd, v[j*D+dglob], acc);
        tmp[idx]=acc*g[idx];
    }
    __syncthreads();
    for(int idx=tid; idx<S*D; idx+=blockDim.x){
        int s=idx>>6, od=idx&63; float acc=ob[od]; const float* tp=&tmp[s*D]; const float* owp=&ow[od*D];
        #pragma unroll
        for(int in=0; in<D; ++in) acc=fmaf(tp[in],owp[in],acc);
        y[((b*S+s)*N+n)*D+od]=acc;
    }
}
torch::Tensor col_attn_forward(torch::Tensor msa_repr, torch::Tensor msa_mask, torch::Tensor ln_w, torch::Tensor ln_b, torch::Tensor qw, torch::Tensor kw, torch::Tensor vw, torch::Tensor gw, torch::Tensor gb, torch::Tensor ow, torch::Tensor ob){ int B=msa_repr.size(0), S=msa_repr.size(1), N=msa_repr.size(2); auto y=torch::empty_like(msa_repr); col_attn_kernel<<<B*N,256>>>(msa_repr.data_ptr<float>(), msa_mask.data_ptr<float>(), ln_w.data_ptr<float>(), ln_b.data_ptr<float>(), qw.data_ptr<float>(), kw.data_ptr<float>(), vw.data_ptr<float>(), gw.data_ptr<float>(), gb.data_ptr<float>(), ow.data_ptr<float>(), ob.data_ptr<float>(), y.data_ptr<float>(), B,S,N); return y; }
'''
_ext = load_inline(name='col_attn_ext_v4unique', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['col_attn_forward'], verbose=False, extra_cuda_cflags=['-O3','--use_fast_math'])
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)