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


_cpp_src = r'''
#include <torch/extension.h>

torch::Tensor forward_cuda(torch::Tensor msa_repr,
                           torch::Tensor msa_mask,
                           torch::Tensor ln_w,
                           torch::Tensor ln_b,
                           torch::Tensor q_w,
                           torch::Tensor k_w,
                           torch::Tensor v_w,
                           torch::Tensor gate_w,
                           torch::Tensor gate_b,
                           torch::Tensor out_w,
                           torch::Tensor out_b,
                           int64_t num_heads);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

#ifndef CUDART_INF_F
#define CUDART_INF_F __int_as_float(0x7f800000)
#endif

__global__ void ln_project_kernel(const float* __restrict__ x,
                                  const float* __restrict__ ln_w,
                                  const float* __restrict__ ln_b,
                                  const float* __restrict__ q_w,
                                  const float* __restrict__ k_w,
                                  const float* __restrict__ v_w,
                                  const float* __restrict__ gate_w,
                                  const float* __restrict__ gate_b,
                                  float* __restrict__ q,
                                  float* __restrict__ k,
                                  float* __restrict__ v,
                                  float* __restrict__ g,
                                  int B, int S, int N, int D, int H, int HD) {
    int token = blockIdx.x; // layout (B,N,S)
    int s = token % S;
    int tmp0 = token / S;
    int n = tmp0 % N;
    int b = tmp0 / N;
    int tid = threadIdx.x;

    __shared__ float red[256];
    __shared__ float xbuf[64];
    __shared__ float mean_s;
    __shared__ float rstd_s;

    float val = 0.0f;
    if (tid < D) {
        val = x[((b * S + s) * N + n) * D + tid];
    }
    red[tid] = (tid < D) ? val : 0.0f;
    __syncthreads();
    for (int off = 128; off > 0; off >>= 1) {
        if (tid < off) red[tid] += red[tid + off];
        __syncthreads();
    }
    if (tid == 0) mean_s = red[0] / (float)D;
    __syncthreads();
    float mean = mean_s;

    float diff = 0.0f;
    if (tid < D) {
        float xv = x[((b * S + s) * N + n) * D + tid];
        diff = xv - mean;
    }
    red[tid] = (tid < D) ? diff * diff : 0.0f;
    __syncthreads();
    for (int off = 128; off > 0; off >>= 1) {
        if (tid < off) red[tid] += red[tid + off];
        __syncthreads();
    }
    if (tid == 0) rstd_s = rsqrtf(red[0] / (float)D + 1.0e-5f);
    __syncthreads();
    float rstd = rstd_s;

    if (tid < D) {
        float xv = x[((b * S + s) * N + n) * D + tid];
        xbuf[tid] = (xv - mean) * rstd * ln_w[tid] + ln_b[tid];
    }
    __syncthreads();

    if (tid < 256) {
        int proj = tid >> 6;       // 0 q, 1 k, 2 v, 3 gate
        int od = tid & 63;
        const float* w = q_w;
        if (proj == 1) w = k_w;
        else if (proj == 2) w = v_w;
        else if (proj == 3) w = gate_w;
        float acc = 0.0f;
        #pragma unroll
        for (int c = 0; c < 64; ++c) {
            acc += xbuf[c] * w[od * D + c];
        }
        if (proj == 3) {
            acc += gate_b[od];
            acc = 1.0f / (1.0f + expf(-acc));
        }
        int hh = od / HD;
        int dd = od - hh * HD;
        int oidx = (((b * N + n) * H + hh) * S + s) * HD + dd;
        if (proj == 0) q[oidx] = acc;
        else if (proj == 1) k[oidx] = acc;
        else if (proj == 2) v[oidx] = acc;
        else g[oidx] = acc;
    }
}

__global__ void attention_kernel(const float* __restrict__ q,
                                 const float* __restrict__ k,
                                 const float* __restrict__ v,
                                 const float* __restrict__ g,
                                 const float* __restrict__ mask,
                                 float* __restrict__ tmp,
                                 int B, int S, int N, int D, int H, int HD) {
    int idx = blockIdx.x; // (B,N,H,S_i)
    int i = idx % S;
    int t0 = idx / S;
    int h = t0 % H;
    t0 /= H;
    int n = t0 % N;
    int b = t0 / N;
    int tid = threadIdx.x; // 0..HD-1

    __shared__ float part[16];
    __shared__ float scores[8];
    __shared__ float weights[8];

    int qi_base = (((b * N + n) * H + h) * S + i) * HD;
    #pragma unroll
    for (int j = 0; j < 8; ++j) {
        if (j < S) {
            int kj_base = (((b * N + n) * H + h) * S + j) * HD;
            float prod = (tid < HD) ? q[qi_base + tid] * k[kj_base + tid] : 0.0f;
            part[tid] = prod;
            __syncthreads();
            if (tid < 8) part[tid] += part[tid + 8];
            __syncthreads();
            if (tid < 4) part[tid] += part[tid + 4];
            __syncthreads();
            if (tid < 2) part[tid] += part[tid + 2];
            __syncthreads();
            if (tid == 0) {
                float m = mask[(b * S + j) * N + n];
                scores[j] = (m == 0.0f) ? -CUDART_INF_F : (part[0] + part[1]) * rsqrtf((float)HD);
            }
            __syncthreads();
        }
    }

    if (tid == 0) {
        float mx = -CUDART_INF_F;
        for (int j = 0; j < S; ++j) mx = fmaxf(mx, scores[j]);
        float sum = 0.0f;
        for (int j = 0; j < S; ++j) {
            float e = expf(scores[j] - mx);
            weights[j] = e;
            sum += e;
        }
        float inv = 1.0f / sum;
        for (int j = 0; j < S; ++j) weights[j] *= inv;
    }
    __syncthreads();

    if (tid < HD) {
        float acc = 0.0f;
        #pragma unroll
        for (int j = 0; j < 8; ++j) {
            if (j < S) {
                int vj_base = (((b * N + n) * H + h) * S + j) * HD;
                acc += weights[j] * v[vj_base + tid];
            }
        }
        acc *= g[qi_base + tid];
        tmp[((b * N + n) * S + i) * D + h * HD + tid] = acc;
    }
}

__global__ void out_proj_kernel(const float* __restrict__ tmp,
                                const float* __restrict__ out_w,
                                const float* __restrict__ out_b,
                                float* __restrict__ out,
                                int B, int S, int N, int D) {
    int token = blockIdx.x; // (B,N,S)
    int s = token % S;
    int t0 = token / S;
    int n = t0 % N;
    int b = t0 / N;
    int tid = threadIdx.x;
    __shared__ float xb[64];
    if (tid < D) xb[tid] = tmp[((b * N + n) * S + s) * D + tid];
    __syncthreads();
    if (tid < D) {
        float acc = out_b[tid];
        #pragma unroll
        for (int c = 0; c < 64; ++c) {
            acc += xb[c] * out_w[tid * D + c];
        }
        out[((b * S + s) * N + n) * D + tid] = acc;
    }
}

torch::Tensor forward_cuda(torch::Tensor msa_repr,
                           torch::Tensor msa_mask,
                           torch::Tensor ln_w,
                           torch::Tensor ln_b,
                           torch::Tensor q_w,
                           torch::Tensor k_w,
                           torch::Tensor v_w,
                           torch::Tensor gate_w,
                           torch::Tensor gate_b,
                           torch::Tensor out_w,
                           torch::Tensor out_b,
                           int64_t num_heads) {
    int B = (int)msa_repr.size(0);
    int S = (int)msa_repr.size(1);
    int N = (int)msa_repr.size(2);
    int D = (int)msa_repr.size(3);
    int H = (int)num_heads;
    int HD = D / H;

    auto opts = msa_repr.options();
    auto q = torch::empty({B, N, H, S, HD}, opts);
    auto k = torch::empty({B, N, H, S, HD}, opts);
    auto v = torch::empty({B, N, H, S, HD}, opts);
    auto g = torch::empty({B, N, H, S, HD}, opts);
    auto tmp = torch::empty({B, N, S, D}, opts);
    auto out = torch::empty_like(msa_repr);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    int tokens = B * N * S;
    ln_project_kernel<<<tokens, 256, 0, stream>>>(
        msa_repr.data_ptr<float>(), ln_w.data_ptr<float>(), ln_b.data_ptr<float>(),
        q_w.data_ptr<float>(), k_w.data_ptr<float>(), v_w.data_ptr<float>(),
        gate_w.data_ptr<float>(), gate_b.data_ptr<float>(),
        q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(), g.data_ptr<float>(),
        B, S, N, D, H, HD);

    attention_kernel<<<B * N * H * S, 16, 0, stream>>>(
        q.data_ptr<float>(), k.data_ptr<float>(), v.data_ptr<float>(), g.data_ptr<float>(),
        msa_mask.data_ptr<float>(), tmp.data_ptr<float>(), B, S, N, D, H, HD);

    out_proj_kernel<<<tokens, 64, 0, stream>>>(
        tmp.data_ptr<float>(), out_w.data_ptr<float>(), out_b.data_ptr<float>(),
        out.data_ptr<float>(), B, S, N, D);
    return out;
}
'''

_ext = load_inline(
    name='col_gated_attn_fused_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['forward_cuda'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3'],
    with_cuda=True,
    verbose=False,
)


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: torch.Tensor, msa_mask: torch.Tensor) -> torch.Tensor:
        return _ext.forward_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,
            int(self.num_heads),
        )