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
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.h>
#include <cuda_runtime.h>
#include <math.h>

__global__ void gated_msa_attention_kernel(
    const float* __restrict__ x,      // (B*N, S, D) after LN view
    const float* __restrict__ Wq,     // (D, D)
    const float* __restrict__ Wk,
    const float* __restrict__ Wv,
    const float* __restrict__ Wg,     // gate (D, D) + bias (D) but bias separate? wait handled in host
    const float* __restrict__ Wo,     // out_proj weight (D, D) + bias later
    const float* __restrict__ bg,     // gate bias (D)
    const float* __restrict__ bo,     // out bias (D)
    float* __restrict__ out,          // final (B, S, N, D) permuted back
    const float* __restrict__ msa_mask_col, // (B, N, S)
    int B, int N, int S, int D, int h, int d,
    float scale
) {
    // Naive single-block-per-column implementation for small sizes
    int idx = blockIdx.x;
    if (idx >= B * N) return;
    int b = idx / N;
    int n = idx % N;

    extern __shared__ float smem[];
    float* Q = smem;
    float* K = Q + S * D;
    float* V = K + S * D;
    float* G = V + S * D;
    float* Attn = G + S * D;
    float* O = Attn + h * S * S;

    // Load x for this column: x shape after perm (B, N, S, D) -> row-major (B*N, S, D)
    const float* x_col = x + (b * N + n) * S * D;

    // Compute Q, K, V, G projections (linear no bias for qkv, with bias for g)
    for (int s = 0; s < S; ++s) {
        for (int dd = threadIdx.x; dd < D; dd += blockDim.x) {
            float qv = 0.0f, kv = 0.0f, vv = 0.0f, gv = 0.0f;
            const float* x_row = x_col + s * D;
            for (int i = 0; i < D; ++i) {
                float xv = x_row[i];
                qv += xv * Wq[i * D + dd];
                kv += xv * Wk[i * D + dd];
                vv += xv * Wv[i * D + dd];
                gv += xv * Wg[i * D + dd];
            }
            gv += bg[dd];
            Q[s * D + dd] = qv;
            K[s * D + dd] = kv;
            V[s * D + dd] = vv;
            G[s * D + dd] = 1.0f / (1.0f + expf(-gv)); // sigmoid
        }
    }
    __syncthreads();

    // Reshape to heads, compute attention per head
    int threads = blockDim.x;
    for (int head = 0; head < h; ++head) {
        for (int i = threadIdx.x; i < S * S; i += threads) {
            int si = i / S;
            int sj = i % S;
            float sum = 0.0f;
            for (int dd = 0; dd < d; ++dd) {
                float qi = Q[si * D + head * d + dd];
                float kj = K[sj * D + head * d + dd];
                sum += qi * kj;
            }
            sum *= scale;
            // mask
            float m = msa_mask_col[b * N * S + n * S + sj];
            Attn[(head * S + si) * S + sj] = (m > 0.5f) ? sum : -1e9f;
        }
    }
    __syncthreads();

    // Softmax per row per head
    for (int head = 0; head < h; ++head) {
        for (int si = threadIdx.x; si < S; si += threads) {
            float maxv = -1e9f;
            for (int sj = 0; sj < S; ++sj) {
                float v = Attn[(head * S + si) * S + sj];
                if (v > maxv) maxv = v;
            }
            float sum = 0.0f;
            for (int sj = 0; sj < S; ++sj) {
                float v = expf(Attn[(head * S + si) * S + sj] - maxv);
                Attn[(head * S + si) * S + sj] = v;
                sum += v;
            }
            for (int sj = 0; sj < S; ++sj) {
                Attn[(head * S + si) * S + sj] /= sum;
            }
        }
    }
    __syncthreads();

    // Compute out = attn @ v , then * g , then reshape/project
    for (int s = 0; s < S; ++s) {
        for (int dd = threadIdx.x; dd < D; dd += threads) {
            int head = dd / d;
            int dh = dd % d;
            float sum = 0.0f;
            for (int sj = 0; sj < S; ++sj) {
                float a = Attn[(head * S + s) * S + sj];
                sum += a * V[sj * D + dd];
            }
            float gated = G[s * D + dd] * sum;
            O[s * D + dd] = gated;
        }
    }
    __syncthreads();

    // Final out_proj: (O @ Wo.T) + bo , then permute back
    float* final_out = out + (b * S + 0) * N * D + n * D;  // stride to (B,S,N,D)
    for (int s = 0; s < S; ++s) {
        for (int dd = threadIdx.x; dd < D; dd += threads) {
            float sum = bo[dd];
            for (int i = 0; i < D; ++i) {
                sum += O[s * D + i] * Wo[i * D + dd];
            }
            // write to (B, S, N, D)[b, s, n, dd]
            final_out[s * N * D + dd] = sum;
        }
    }
}

torch::Tensor gated_msa_cuda(
    torch::Tensor msa_repr,
    torch::Tensor msa_mask,
    torch::Tensor ln_weight,
    torch::Tensor ln_bias,
    torch::Tensor Wq,
    torch::Tensor Wk,
    torch::Tensor Wv,
    torch::Tensor Wg,
    torch::Tensor bg,
    torch::Tensor Wo,
    torch::Tensor bo,
    int num_heads
) {
    auto B = msa_repr.size(0);
    auto S = msa_repr.size(1);
    auto N = msa_repr.size(2);
    auto D = msa_repr.size(3);
    int h = num_heads;
    int d = D / h;
    float scale = 1.0f / sqrtf((float)d);

    // LayerNorm on permuted (B, N, S, D)
    auto x = msa_repr.permute({0,2,1,3}).contiguous(); // (B,N,S,D)
    auto ln_mean = x.mean(-1, true);
    auto ln_var = x.var(-1, true, false);
    auto x_ln = (x - ln_mean) / torch::sqrt(ln_var + 1e-5);
    x_ln = x_ln * ln_weight.view({1,1,1,D}) + ln_bias.view({1,1,1,D});
    x_ln = x_ln.contiguous();

    auto msa_mask_col = msa_mask.permute({0,2,1}).contiguous(); // (B,N,S)

    auto output = torch::zeros({B, S, N, D}, msa_repr.options());

    int blocks = B * N;
    int threads = 256;
    size_t smem = (4 * S * D + h * S * S + S * D) * sizeof(float); // QKV G O + Attn

    gated_msa_attention_kernel<<<blocks, threads, smem>>>(
        x_ln.data_ptr<float>(),
        Wq.data_ptr<float>(),
        Wk.data_ptr<float>(),
        Wv.data_ptr<float>(),
        Wg.data_ptr<float>(),
        Wo.data_ptr<float>(),
        bg.data_ptr<float>(),
        bo.data_ptr<float>(),
        output.data_ptr<float>(),
        msa_mask_col.data_ptr<float>(),
        B, N, S, D, h, d, scale
    );

    return output;
}
'''

cpp_source = r'''
torch::Tensor gated_msa_cuda(
    torch::Tensor msa_repr,
    torch::Tensor msa_mask,
    torch::Tensor ln_weight,
    torch::Tensor ln_bias,
    torch::Tensor Wq,
    torch::Tensor Wk,
    torch::Tensor Wv,
    torch::Tensor Wg,
    torch::Tensor bg,
    torch::Tensor Wo,
    torch::Tensor bo,
    int num_heads
);
'''

ext = load_inline(
    name='gated_msa_ext',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['gated_msa_cuda'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

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)  # dropout=0 in inference

    def forward(self, msa_repr: torch.Tensor, msa_mask: torch.Tensor) -> torch.Tensor:
        B, S, N, D = msa_repr.shape
        // Delegate to custom CUDA
        out = ext.gated_msa_cuda(
            msa_repr,
            msa_mask,
            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,
            self.num_heads
        )
        return out
```