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
import torch.nn.functional as F
import math
from torch.utils.cpp_extension import load_inline

cuda_source = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

__global__ void attention_kernel(
    const float* q, const float* k, const float* v, const float* bias, const float* mask,
    float* out, float* gate,
    int B, int S, int h, int N, int d, float scale
) {
    // Very naive implementation for small N=32
    int bsh = blockIdx.x;
    int b = bsh / (S * h);
    int sh = bsh % (S * h);
    int s = sh / h;
    int head = sh % h;
    
    extern __shared__ float smem[];
    
    // Compute attn scores
    for(int i = threadIdx.x; i < N; i += blockDim.x) {
        float sum = 0.0f;
        for(int j=0; j<N; j++) {
            float score = 0.0f;
            for(int dd=0; dd<d; dd++) {
                score += q[((b*S + s)*h + head)*N*d + i*d + dd] * k[((b*S + s)*h + head)*N*d + j*d + dd];
            }
            score = score * scale + bias[(b*h + head)*N*N + i*N + j];
            if (mask[(b*S + s)*N + j] == 0) score = -INFINITY;
            smem[j] = score;
        }
        // softmax
        float maxv = -INFINITY;
        for(int j=0; j<N; j++) maxv = fmaxf(maxv, smem[j]);
        float sumexp = 0.0f;
        for(int j=0; j<N; j++) {
            smem[j] = expf(smem[j] - maxv);
            sumexp += smem[j];
        }
        for(int j=0; j<N; j++) smem[j] /= sumexp;
        
        // matmul with v and gate
        for(int dd=0; dd<d; dd++) {
            float res = 0.0f;
            for(int j=0; j<N; j++) {
                res += smem[j] * v[((b*S + s)*h + head)*N*d + j*d + dd];
            }
            int out_idx = ((b*S + s)*h + head)*N*d + i*d + dd;
            out[out_idx] = res * gate[out_idx];
        }
    }
}

torch::Tensor custom_attention(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor bias, torch::Tensor g, torch::Tensor mask, float scale) {
    auto options = torch::TensorOptions().dtype(torch::kFloat32).device(q.device());
    auto out = torch::empty_like(q);
    int B = q.size(0);
    int S = q.size(1);
    int h = q.size(2);
    int N = q.size(3);
    int d = q.size(4);
    dim3 grid(B*S*h);
    int threads = 32;
    attention_kernel<<<grid, threads, N*sizeof(float)>>>(q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(), bias.data_ptr<float>(), mask.data_ptr<float>(), out.data_ptr<float>(), g.data_ptr<float>(), B, S, h, N, d, scale);
    return out;
}
'''

cpp_source = r'''
torch::Tensor custom_attention(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor bias, torch::Tensor g, torch::Tensor mask, float scale);
'''

ext = load_inline(
    name='custom_attn',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['custom_attention'],
    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)

    def forward(self, msa_repr, pair_repr, msa_mask):
        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).contiguous()
        k = self.key(x).view(B, S, N, h, d).permute(0, 1, 3, 2, 4).contiguous()
        v = self.value(x).view(B, S, N, h, d).permute(0, 1, 3, 2, 4).contiguous()
        g = torch.sigmoid(self.gate(x)).view(B, S, N, h, d).permute(0, 1, 3, 2, 4).contiguous()
        bias = self.pair_bias(pair_repr).permute(0, 3, 1, 2).contiguous()
        scale = 1.0 / math.sqrt(d)
        out = ext.custom_attention(q, k, v, bias, g, msa_mask, scale)
        out = out.permute(0, 1, 3, 2, 4).reshape(B, S, N, D)
        return self.out_proj(out)