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_src = r'''
#include <torch/extension.h>

torch::Tensor row_msa_attn_forward(
    torch::Tensor msa, torch::Tensor pair, torch::Tensor 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 pbw,
    torch::Tensor ow, torch::Tensor ob,
    int64_t num_heads);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math_constants.h>

__global__ void prep_qkvg_kernel(
    const float* __restrict__ msa,
    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,
    float* __restrict__ q,
    float* __restrict__ k,
    float* __restrict__ v,
    float* __restrict__ g,
    int T, int D)
{
    int tok = blockIdx.x;
    int tid = threadIdx.x;
    __shared__ float red[128];
    __shared__ float xn[128];

    float s = 0.0f;
    for (int c = tid; c < D; c += blockDim.x) s += msa[tok * D + c];
    red[tid] = s;
    __syncthreads();
    for (int off = blockDim.x >> 1; off > 0; off >>= 1) {
        if (tid < off) red[tid] += red[tid + off];
        __syncthreads();
    }
    float mean = red[0] / (float)D;

    float vsum = 0.0f;
    for (int c = tid; c < D; c += blockDim.x) {
        float z = msa[tok * D + c] - mean;
        vsum += z * z;
    }
    red[tid] = vsum;
    __syncthreads();
    for (int off = blockDim.x >> 1; off > 0; off >>= 1) {
        if (tid < off) red[tid] += red[tid + off];
        __syncthreads();
    }
    float invstd = rsqrtf(red[0] / (float)D + 1.0e-5f);

    if (tid < D) {
        float y = (msa[tok * D + tid] - mean) * invstd;
        xn[tid] = y * ln_w[tid] + ln_b[tid];
    }
    __syncthreads();

    if (tid < D) {
        int oc = tid;
        float sq = 0.0f, sk = 0.0f, sv = 0.0f, sg = 0.0f;
        const float* wq = qw + oc * D;
        const float* wk = kw + oc * D;
        const float* wv = vw + oc * D;
        const float* wg = gw + oc * D;
        #pragma unroll
        for (int c = 0; c < 64; ++c) {
            if (c < D) {
                float x = xn[c];
                sq += x * wq[c];
                sk += x * wk[c];
                sv += x * wv[c];
                sg += x * wg[c];
            }
        }
        q[tok * D + oc] = sq;
        k[tok * D + oc] = sk;
        v[tok * D + oc] = sv;
        float zz = sg + gb[oc];
        g[tok * D + oc] = 1.0f / (1.0f + expf(-zz));
    }
}

__global__ void pair_bias_kernel(
    const float* __restrict__ pair,
    const float* __restrict__ pbw,
    float* __restrict__ pb,
    int B, int N, int PD, int H)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * H * N * N;
    if (idx >= total) return;
    int j = idx % N;
    int t = idx / N;
    int i = t % N;
    t /= N;
    int h = t % H;
    int b = t / H;
    const float* pr = pair + ((b * N + i) * N + j) * PD;
    const float* w = pbw + h * PD;
    float sum = 0.0f;
    #pragma unroll
    for (int p = 0; p < 64; ++p) {
        if (p < PD) sum += pr[p] * w[p];
    }
    pb[idx] = sum;
}

__global__ void attn_kernel(
    const float* __restrict__ q,
    const float* __restrict__ k,
    const float* __restrict__ v,
    const float* __restrict__ gate,
    const float* __restrict__ pb,
    const float* __restrict__ mask,
    float* __restrict__ mid,
    int B, int S, int N, int D, int H, int d)
{
    int bid = blockIdx.x;
    int tid = threadIdx.x;
    int i = bid % N;
    int t = bid / N;
    int h = t % H;
    t /= H;
    int s = t % S;
    int b = t / S;

    __shared__ float vals[64];
    __shared__ float red[64];

    int tok_i = (b * S + s) * N + i;
    int qbase = tok_i * D + h * d;
    float score = -CUDART_INF_F;
    if (tid < N) {
        int j = tid;
        int tok_j = (b * S + s) * N + j;
        int kbase = tok_j * D + h * d;
        float dot = 0.0f;
        #pragma unroll
        for (int r = 0; r < 16; ++r) {
            if (r < d) dot += q[qbase + r] * k[kbase + r];
        }
        score = dot * rsqrtf((float)d) + pb[((b * H + h) * N + i) * N + j];
        if (mask[(b * S + s) * N + j] == 0.0f) score = -CUDART_INF_F;
        vals[j] = score;
    }
    red[tid] = (tid < N) ? score : -CUDART_INF_F;
    __syncthreads();
    for (int off = 32; off > 0; off >>= 1) {
        if (tid < off) red[tid] = fmaxf(red[tid], red[tid + off]);
        __syncthreads();
    }
    float m = red[0];
    float e = 0.0f;
    if (tid < N) {
        e = expf(vals[tid] - m);
        vals[tid] = e;
    }
    red[tid] = (tid < N) ? e : 0.0f;
    __syncthreads();
    for (int off = 32; off > 0; off >>= 1) {
        if (tid < off) red[tid] += red[tid + off];
        __syncthreads();
    }
    float invsum = 1.0f / red[0];

    if (tid < d) {
        int r = tid;
        float out = 0.0f;
        for (int j = 0; j < N; ++j) {
            int tok_j = (b * S + s) * N + j;
            out += vals[j] * invsum * v[tok_j * D + h * d + r];
        }
        mid[qbase + r] = out * gate[qbase + r];
    }
}

__global__ void out_proj_kernel(
    const float* __restrict__ mid,
    const float* __restrict__ ow,
    const float* __restrict__ ob,
    float* __restrict__ out,
    int T, int D)
{
    int tok = blockIdx.x;
    int tid = threadIdx.x;
    if (tok >= T) return;
    if (tid < D) {
        int oc = tid;
        const float* w = ow + oc * D;
        float sum = ob[oc];
        #pragma unroll
        for (int c = 0; c < 64; ++c) {
            if (c < D) sum += mid[tok * D + c] * w[c];
        }
        out[tok * D + oc] = sum;
    }
}

torch::Tensor row_msa_attn_forward(
    torch::Tensor msa, torch::Tensor pair, torch::Tensor 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 pbw,
    torch::Tensor ow, torch::Tensor ob,
    int64_t num_heads)
{
    int B = (int)msa.size(0);
    int S = (int)msa.size(1);
    int N = (int)msa.size(2);
    int D = (int)msa.size(3);
    int PD = (int)pair.size(3);
    int H = (int)num_heads;
    int d = D / H;
    int T = B * S * N;

    auto q = torch::empty({T, D}, msa.options());
    auto k = torch::empty({T, D}, msa.options());
    auto v = torch::empty({T, D}, msa.options());
    auto g = torch::empty({T, D}, msa.options());
    auto pb = torch::empty({B, H, N, N}, msa.options());
    auto mid = torch::empty({T, D}, msa.options());
    auto out = torch::empty({B, S, N, D}, msa.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    prep_qkvg_kernel<<<T, 128, 0, stream>>>(
        msa.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>(),
        q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(), g.data_ptr<float>(), T, D);

    int pbtot = B * H * N * N;
    pair_bias_kernel<<<(pbtot + 127) / 128, 128, 0, stream>>>(
        pair.data_ptr<float>(), pbw.data_ptr<float>(), pb.data_ptr<float>(), B, N, PD, H);

    attn_kernel<<<B * S * H * N, 64, 0, stream>>>(
        q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(), g.data_ptr<float>(),
        pb.data_ptr<float>(), mask.data_ptr<float>(), mid.data_ptr<float>(),
        B, S, N, D, H, d);

    out_proj_kernel<<<T, 128, 0, stream>>>(
        mid.data_ptr<float>(), ow.data_ptr<float>(), ob.data_ptr<float>(), out.data_ptr<float>(), T, D);

    return out;
}
'''

_ext = load_inline(
    name="row_msa_attn_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["row_msa_attn_forward"],
    extra_cuda_cflags=["-O3"],
    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
        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:
        return _ext.row_msa_attn_forward(
            msa_repr.contiguous(), pair_repr.contiguous(), msa_mask.contiguous(),
            self.layer_norm.weight.contiguous(), self.layer_norm.bias.contiguous(),
            self.query.weight.contiguous(), self.key.weight.contiguous(), self.value.weight.contiguous(),
            self.gate.weight.contiguous(), self.gate.bias.contiguous(),
            self.pair_bias.weight.contiguous(),
            self.out_proj.weight.contiguous(), self.out_proj.bias.contiguous(),
            int(self.num_heads),
        )