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

# Define CUDA source for fused operations
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <c10/cuda/CUDAGuard.h>
#include <ATen/cuda/CUDAContext.h>

template <typename scalar_t>
__global__ void layer_norm_kernel(
    const scalar_t* __restrict__ input,
    scalar_t* __restrict__ output,
    const scalar_t* __restrict__ weight,
    const scalar_t* __restrict__ bias,
    int N, int D, float eps) {
    
    int idx = blockIdx.x;
    int tid = threadIdx.x;
    
    if (idx >= N) return;
    
    __shared__ float mean;
    __shared__ float var;
    __shared__ float inv_std;
    
    float sum = 0.0f;
    float sum_sq = 0.0f;
    
    for (int i = tid; i < D; i += blockDim.x) {
        float val = (float)input[idx * D + i];
        sum += val;
        sum_sq += val * val;
    }
    
    // Block reduction for sum
    for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
        sum += __shfl_down_sync(0xffffffff, sum, offset);
        sum_sq += __shfl_down_sync(0xffffffff, sum_sq, offset);
    }
    
    if (tid == 0) {
        mean = sum / D;
        var = sum_sq / D - mean * mean;
        inv_std = 1.0f / sqrtf(var + eps);
    }
    __syncthreads();
    
    for (int i = tid; i < D; i += blockDim.x) {
        float val = (float)input[idx * D + i];
        val = (val - mean) * inv_std;
        if (weight != nullptr) val *= (float)weight[i];
        if (bias != nullptr) val += (float)bias[i];
        output[idx * D + i] = (scalar_t)val;
    }
}

template <typename scalar_t>
__global__ void fused_gate_out_kernel(
    const scalar_t* __restrict__ gate_proj,
    const scalar_t* __restrict__ attn_out,
    const scalar_t* __restrict__ out_weight,
    const scalar_t* __restrict__ out_bias,
    scalar_t* __restrict__ output,
    int B, int N, int S, int D, int h, int d) {
    
    // Each block handles one (b, n, s) position
    int b = blockIdx.x / (N * S);
    int ns = blockIdx.x % (N * S);
    int n = ns / S;
    int s = ns % S;
    int tid = threadIdx.x;
    
    extern __shared__ float shared_mem[];
    float* local_out = shared_mem;  // Size D
    float* local_gate = &shared_mem[D];  // Size D
    
    // Initialize output
    for (int i = tid; i < D; i += blockDim.x) {
        local_out[i] = 0.0f;
        local_gate[i] = 0.0f;
    }
    __syncthreads();
    
    // Aggregate attention output across heads and compute sigmoid gate
    for (int hd = tid; hd < h * d; hd += blockDim.x) {
        int head = hd / d;
        int dim = hd % d;
        
        // attn_out shape: (B, N, h, S, d) -> index: b*N*h*S*d + n*h*S*d + head*S*d + s*d + dim
        int attn_idx = ((b * N + n) * h + head) * S * d + s * d + dim;
        float val = (float)attn_out[attn_idx];
        
        // gate_proj shape: (B, N, S, D) before sigmoid
        int gate_idx = (b * N + n) * S * D + s * D + head * d + dim;
        float gate_val = (float)gate_proj[gate_idx];
        
        atomicAdd(&local_gate[head * d + dim], gate_val);
        atomicAdd(&local_out[head * d + dim], val);
    }
    __syncthreads();
    
    // Apply sigmoid to gate and multiply
    for (int i = tid; i < D; i += blockDim.x) {
        float gate = 1.0f / (1.0f + expf(-local_gate[i]));
        local_out[i] = gate * local_out[i];
    }
    __syncthreads();
    
    // Output projection: local_out (D) @ out_weight (D, D) + bias
    for (int i = tid; i < D; i += blockDim.x) {
        float sum = (out_bias != nullptr) ? (float)out_bias[i] : 0.0f;
        for (int j = 0; j < D; j++) {
            sum += local_out[j] * (float)out_weight[j * D + i];
        }
        // Write to output: (B, S, N, D)
        // Input was (B, N, S, D), need to permute back
        int out_idx = b * S * N * D + s * N * D + n * D + i;
        output[out_idx] = (scalar_t)sum;
    }
}

template <typename scalar_t>
__global__ void prepare_qkv_kernel(
    const scalar_t* __restrict__ ln_out,
    const scalar_t* __restrict__ q_weight,
    const scalar_t* __restrict__ k_weight,
    const scalar_t* __restrict__ v_weight,
    const scalar_t* __restrict__ g_weight,
    const scalar_t* __restrict__ g_bias,
    scalar_t* __restrict__ q_out,
    scalar_t* __restrict__ k_out,
    scalar_t* __restrict__ v_out,
    scalar_t* __restrict__ g_out,
    int B, int N, int S, int D, int h, int d) {
    
    // Each block handles one (b, n, s) position
    int idx = blockIdx.x;
    int tid = threadIdx.x;
    
    if (idx >= B * N * S) return;
    
    int b = idx / (N * S);
    int ns = idx % (N * S);
    int n = ns / S;
    int s = ns % S;
    
    extern __shared__ float shared_mem[];
    float* ln_local = shared_mem;  // Size D
    
    // Load layer norm output
    for (int i = tid; i < D; i += blockDim.x) {
        ln_local[i] = (float)ln_out[idx * D + i];
    }
    __syncthreads();
    
    // Compute Q, K, V, G projections and write in (B, N, h, S, d) format
    for (int hd = tid; hd < h * d; hd += blockDim.x) {
        int head = hd / d;
        int dim = hd % d;
        
        float q_val = 0.0f, k_val = 0.0f, v_val = 0.0f, g_val = (g_bias != nullptr) ? (float)g_bias[head * d + dim] : 0.0f;
        
        for (int j = 0; j < D; j++) {
            float ln_val = ln_local[j];
            q_val += ln_val * (float)q_weight[j * h * d + head * d + dim];
            k_val += ln_val * (float)k_weight[j * h * d + head * d + dim];
            v_val += ln_val * (float)v_weight[j * h * d + head * d + dim];
            g_val += ln_val * (float)g_weight[j * h * d + head * d + dim];
        }
        
        // Write Q, K, V in (B, N, h, S, d) format
        int out_idx = ((b * N + n) * h + head) * S * d + s * d + dim;
        q_out[out_idx] = (scalar_t)q_val;
        k_out[out_idx] = (scalar_t)k_val;
        v_out[out_idx] = (scalar_t)v_val;
        g_out[out_idx] = (scalar_t)g_val;
    }
}

template <typename scalar_t>
__global__ void softmax_mask_kernel(
    scalar_t* __restrict__ attn,
    const int* __restrict__ mask,
    int B, int N, int h, int S) {
    
    // Each block handles one (b, n, head, i) row
    int idx = blockIdx.x;
    int tid = threadIdx.x;
    
    if (idx >= B * N * h * S) return;
    
    int b = idx / (N * h * S);
    int remainder = idx % (N * h * S);
    int n = remainder / (h * S);
    remainder = remainder % (h * S);
    int head = remainder / S;
    int i = remainder % S;
    
    // Find mask index: (B, N, S) -> mask[b, n, j] for all j
    int mask_base = b * N * S + n * S;
    
    __shared__ float max_val;
    __shared__ float sum_exp;
    
    float local_max = -1e30f;
    
    // Apply mask and find max
    for (int j = tid; j < S; j += blockDim.x) {
        int attn_idx = ((b * N + n) * h + head) * S * S + i * S + j;
        if (mask[mask_base + j] == 0) {
            attn[attn_idx] = (scalar_t)(-1e30f);
        }
        float val = (float)attn[attn_idx];
        local_max = fmaxf(local_max, val);
    }
    
    // Block reduction for max
    for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
        local_max = fmaxf(local_max, __shfl_down_sync(0xffffffff, local_max, offset));
    }
    
    if (tid == 0) {
        max_val = local_max;
    }
    __syncthreads();
    
    // Compute exp and sum
    float local_sum = 0.0f;
    for (int j = tid; j < S; j += blockDim.x) {
        int attn_idx = ((b * N + n) * h + head) * S * S + i * S + j;
        float val = expf((float)attn[attn_idx] - max_val);
        attn[attn_idx] = (scalar_t)val;
        local_sum += val;
    }
    
    // Block reduction for sum
    for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
        local_sum += __shfl_down_sync(0xffffffff, local_sum, offset);
    }
    
    if (tid == 0) {
        sum_exp = local_sum;
    }
    __syncthreads();
    
    // Normalize
    float inv_sum = 1.0f / sum_exp;
    for (int j = tid; j < S; j += blockDim.x) {
        int attn_idx = ((b * N + n) * h + head) * S * S + i * S + j;
        attn[attn_idx] = (scalar_t)((float)attn[attn_idx] * inv_sum);
    }
}

std::vector<torch::Tensor> msa_attention_cuda(
    torch::Tensor msa_repr,
    torch::Tensor msa_mask,
    torch::Tensor ln_weight,
    torch::Tensor ln_bias,
    torch::Tensor q_weight,
    torch::Tensor k_weight,
    torch::Tensor v_weight,
    torch::Tensor g_weight,
    torch::Tensor g_bias,
    torch::Tensor out_weight,
    torch::Tensor out_bias,
    int num_heads,
    float scale) {
    
    auto stream = at::cuda::getCurrentCUDAStream();
    
    int B = msa_repr.size(0);
    int S = msa_repr.size(1);
    int N = msa_repr.size(2);
    int D = msa_repr.size(3);
    int h = num_heads;
    int d = D / h;
    
    // Permute input: (B, S, N, D) -> (B, N, S, D)
    auto msa_permuted = msa_repr.permute({0, 2, 1, 3}).contiguous();
    
    // Layer norm output
    auto ln_out = torch::empty_like(msa_permuted);
    
    // Apply layer norm
    int ln_threads = std::min(D, 256);
    int ln_blocks = B * N * S;
    
    AT_DISPATCH_FLOATING_TYPES(msa_permuted.scalar_type(), "layer_norm_kernel", ([&] {
        layer_norm_kernel<scalar_t><<<ln_blocks, ln_threads, 0, stream>>>(
            msa_permuted.data_ptr<scalar_t>(),
            ln_out.data_ptr<scalar_t>(),
            ln_weight.data_ptr<scalar_t>(),
            ln_bias.data_ptr<scalar_t>(),
            B * N * S, D, 1e-5f);
    }));
    
    // Prepare Q, K, V, G
    auto q_out = torch::empty({B, N, h, S, d}, msa_repr.options());
    auto k_out = torch::empty({B, N, h, S, d}, msa_repr.options());
    auto v_out = torch::empty({B, N, h, S, d}, msa_repr.options());
    auto g_proj = torch::empty({B, N, S, D}, msa_repr.options());
    
    int prep_threads = std::min(h * d, 256);
    int prep_blocks = B * N * S;
    size_t prep_shared = D * sizeof(float);
    
    AT_DISPATCH_FLOATING_TYPES(msa_permuted.scalar_type(), "prepare_qkv_kernel", ([&] {
        prepare_qkv_kernel<scalar_t><<<prep_blocks, prep_threads, prep_shared, stream>>>(
            ln_out.data_ptr<scalar_t>(),
            q_weight.data_ptr<scalar_t>(),
            k_weight.data_ptr<scalar_t>(),
            v_weight.data_ptr<scalar_t>(),
            g_weight.data_ptr<scalar_t>(),
            g_bias.data_ptr<scalar_t>(),
            q_out.data_ptr<scalar_t>(),
            k_out.data_ptr<scalar_t>(),
            v_out.data_ptr<scalar_t>(),
            g_proj.data_ptr<scalar_t>(),
            B, N, S, D, h, d);
    }));
    
    // Compute attention: Q @ K^T
    // q_out: (B, N, h, S, d), k_out: (B, N, h, S, d)
    // Result: (B, N, h, S, S)
    
    // Reshape for batched matmul: (B*N*h, S, d)
    auto q_reshaped = q_out.view({B * N * h, S, d});
    auto k_reshaped = k_out.view({B * N * h, S, d});
    
    // attention = q @ k^T
    auto attn = torch::matmul(q_reshaped, k_reshaped.transpose(1, 2));
    attn = attn.view({B, N, h, S, S});
    
    // Apply scale
    attn = attn * scale;
    
    // Apply mask and softmax
    // msa_mask: (B, S, N) -> col_mask: (B, N, S)
    auto col_mask = msa_mask.permute({0, 2, 1}).contiguous();
    
    int softmax_threads = std::min(S, 256);
    int softmax_blocks = B * N * h * S;
    
    AT_DISPATCH_FLOATING_TYPES(attn.scalar_type(), "softmax_mask_kernel", ([&] {
        softmax_mask_kernel<scalar_t><<<softmax_blocks, softmax_threads, 0, stream>>>(
            attn.data_ptr<scalar_t>(),
            col_mask.to(at::kInt).data_ptr<int>(),
            B, N, h, S);
    }));
    
    // Apply attention to values: attn @ v
    // attn: (B, N, h, S, S), v: (B, N, h, S, d)
    auto attn_reshaped = attn.view({B * N * h, S, S});
    auto v_reshaped = v_out.view({B * N * h, S, d});
    
    auto attn_out = torch::matmul(attn_reshaped, v_reshaped);
    attn_out = attn_out.view({B, N, h, S, d});
    
    // Fused gate + output projection
    auto output = torch::empty({B, S, N, D}, msa_repr.options());
    
    int fused_threads = std::min(D, 256);
    int fused_blocks = B * N * S;
    size_t fused_shared = 2 * D * sizeof(float);
    
    AT_DISPATCH_FLOATING_TYPES(msa_permuted.scalar_type(), "fused_gate_out_kernel", ([&] {
        fused_gate_out_kernel<scalar_t><<<fused_blocks, fused_threads, fused_shared, stream>>>(
            g_proj.data_ptr<scalar_t>(),
            attn_out.data_ptr<scalar_t>(),
            out_weight.data_ptr<scalar_t>(),
            out_bias.defined() ? out_bias.data_ptr<scalar_t>() : nullptr,
            output.data_ptr<scalar_t>(),
            B, N, S, D, h, d);
    }));
    
    return {output};
}
"""

# Load the CUDA extension
try:
    msa_attention_ext = load_inline(
        name='msa_attention_ext',
        cpp_sources='',
        cuda_sources=cuda_source,
        functions=['msa_attention_cuda'],
        verbose=True,
        extra_cuda_cflags=['-O3', '--use_fast_math']
    )
except Exception as e:
    print(f"Failed to load CUDA extension: {e}")
    msa_attention_ext = None


class ModelNew(nn.Module):
    """
    Column-wise gated self-attention over MSA representations with CUDA kernels.
    """

    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)
        
        self.msa_dim = msa_dim

    def forward(self, msa_repr: torch.Tensor, msa_mask: torch.Tensor) -> torch.Tensor:
        B, S, N, D = msa_repr.shape
        h = self.num_heads
        d = self.head_dim
        
        # Use custom CUDA kernel if available and on CUDA
        if msa_attention_ext is not None and msa_repr.is_cuda:
            scale = 1.0 / math.sqrt(d)
            result = msa_attention_ext.msa_attention_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,
                scale
            )
            out = result[0]
            # Apply dropout (only during training)
            if self.training and self.dropout.p > 0:
                out = self.dropout(out)
            return out
        
        # Fallback to PyTorch implementation
        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)