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

# C++ source with function declaration
cpp_source = """
torch::Tensor fused_attention_cuda(
    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor g,
    torch::Tensor bias, c10::optional<torch::Tensor> msa_mask,
    int64_t B, int64_t S, int64_t N, int64_t h, int64_t d, double scale);
"""

# CUDA source with kernel implementation
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

template <typename scalar_t, int BLOCK_SIZE>
__global__ void fused_attention_kernel(
    const scalar_t* __restrict__ q,
    const scalar_t* __restrict__ k,
    const scalar_t* __restrict__ v,
    const scalar_t* __restrict__ g,
    const scalar_t* __restrict__ bias,
    const bool* __restrict__ mask,
    scalar_t* __restrict__ out,
    int B, int S, int N, int h, int d,
    float scale) {
    
    // Each block handles one (batch, seq, head, query_pos)
    // Tensor layout: (B, S, h, N, d)
    int batch_idx = blockIdx.x / (S * h);
    int s_idx = (blockIdx.x / h) % S;
    int head_idx = blockIdx.x % h;
    int query_pos = blockIdx.y;
    
    if (batch_idx >= B || s_idx >= S || head_idx >= h || query_pos >= N) return;
    
    int tid = threadIdx.x;
    
    // Base offset for this (batch, seq, head)
    // Index: b * S * h * N * d + s * h * N * d + head * N * d + n * d + i
    int base_offset = (batch_idx * S + s_idx) * h * N * d + head_idx * N * d;
    int q_offset = base_offset + query_pos * d;  // Position for Q
    
    // Shared memory for Q (all threads need access to full Q vector)
    __shared__ float Q_shared[64];  // max d=64
    
    // Load Q into shared memory
    for (int i = tid; i < d; i += BLOCK_SIZE) {
        Q_shared[i] = float(q[q_offset + i]);
    }
    __syncthreads();
    
    // Shared memory for reduction
    __shared__ float max_score_shared;
    __shared__ float sum_exp_shared;
    __shared__ float out_shared[64];  // max d
    
    if (tid == 0) {
        max_score_shared = -INFINITY;
        sum_exp_shared = 0.0f;
    }
    for (int i = tid; i < d; i += BLOCK_SIZE) {
        out_shared[i] = 0.0f;
    }
    __syncthreads();
    
    // First pass: find max score
    float local_max = -INFINITY;
    for (int kv_pos = tid; kv_pos < N; kv_pos += BLOCK_SIZE) {
        int mask_idx = batch_idx * S * N + s_idx * N + kv_pos;
        if (mask != nullptr && !mask[mask_idx]) continue;
        
        int k_offset = base_offset + kv_pos * d;  // K position
        
        float score = 0.0f;
        for (int i = 0; i < d; i++) {
            score += Q_shared[i] * float(k[k_offset + i]);
        }
        score *= scale;
        
        // Bias: (B, h, N, N) -> bias[b * h * N * N + head * N * N + q * N + kv]
        int bias_idx = batch_idx * h * N * N + head_idx * N * N + query_pos * N + kv_pos;
        score += float(bias[bias_idx]);
        
        local_max = max(local_max, score);
    }
    
    // Reduce max within warp
    for (int offset = 16; offset > 0; offset >>= 1) {
        local_max = max(local_max, __shfl_down_sync(0xffffffff, local_max, offset));
    }
    
    if ((tid & 31) == 0) {
        atomicMax(reinterpret_cast<int*>(&max_score_shared), __float_as_int(local_max));
    }
    __syncthreads();
    float global_max = max_score_shared;
    
    // Second pass: compute exp and accumulate
    float local_sum = 0.0f;
    float local_out[64];
    for (int i = 0; i < d; i++) local_out[i] = 0.0f;
    
    for (int kv_pos = tid; kv_pos < N; kv_pos += BLOCK_SIZE) {
        int mask_idx = batch_idx * S * N + s_idx * N + kv_pos;
        bool valid = (mask == nullptr) || mask[mask_idx];
        
        if (valid) {
            int k_offset = base_offset + kv_pos * d;
            
            float score = 0.0f;
            for (int i = 0; i < d; i++) {
                score += Q_shared[i] * float(k[k_offset + i]);
            }
            score *= scale;
            
            int bias_idx = batch_idx * h * N * N + head_idx * N * N + query_pos * N + kv_pos;
            score += float(bias[bias_idx]);
            
            float exp_val = expf(score - global_max);
            local_sum += exp_val;
            
            int v_offset = base_offset + kv_pos * d;  // V position
            for (int i = 0; i < d; i++) {
                local_out[i] += exp_val * float(v[v_offset + i]);
            }
        }
    }
    
    // Reduce sum within warp
    for (int offset = 16; offset > 0; offset >>= 1) {
        local_sum += __shfl_down_sync(0xffffffff, local_sum, offset);
    }
    
    if ((tid & 31) == 0) {
        atomicAdd(&sum_exp_shared, local_sum);
    }
    
    // Reduce out within warp (each thread handles subset of d)
    for (int i = 0; i < d; i++) {
        float val = local_out[i];
        for (int offset = 16; offset > 0; offset >>= 1) {
            val += __shfl_down_sync(0xffffffff, val, offset);
        }
        if ((tid & 31) == 0) {
            atomicAdd(&out_shared[i], val);
        }
    }
    __syncthreads();
    
    // Write output with gate
    // Output layout: (B, S, h, N, d) - same as Q, K, V, G
    int out_offset = base_offset + query_pos * d;
    
    float inv_sum = 1.0f / (sum_exp_shared + 1e-8f);
    for (int i = tid; i < d; i += BLOCK_SIZE) {
        float norm_out = out_shared[i] * inv_sum;
        float gate_val = float(g[out_offset + i]);
        out[out_offset + i] = scalar_t(gate_val * norm_out);
    }
}

torch::Tensor fused_attention_cuda(
    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor g,
    torch::Tensor bias, c10::optional<torch::Tensor> msa_mask,
    int64_t B, int64_t S, int64_t N, int64_t h, int64_t d, double scale) {
    
    auto options = q.options();
    // Output is (B, S, h, N, d) - same layout as input tensors
    auto out = torch::empty({B, S, int(h), int(N), int(d)}, options);
    
    const int threads = 128;
    
    dim3 grid(B * S * h, N);
    dim3 block(threads);
    
    bool* mask_ptr = msa_mask.has_value() ? msa_mask.value().data_ptr<bool>() : nullptr;
    
    AT_DISPATCH_FLOATING_TYPES(q.scalar_type(), "fused_attention_kernel", ([&] {
        fused_attention_kernel<scalar_t, threads><<<grid, block>>>(
            q.data_ptr<scalar_t>(),
            k.data_ptr<scalar_t>(),
            v.data_ptr<scalar_t>(),
            g.data_ptr<scalar_t>(),
            bias.data_ptr<scalar_t>(),
            mask_ptr,
            out.data_ptr<scalar_t>(),
            B, S, N, h, d,
            static_cast<float>(scale)
        );
    }));
    
    return out;
}
"""

# Load the CUDA extension
fused_ops = load_inline(
    name='fused_ops_v6',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['fused_attention_cuda'],
    extra_cuda_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)
        self.dropout_p = dropout

    def forward(
        self,
        msa_repr: torch.Tensor,
        pair_repr: torch.Tensor,
        msa_mask: torch.Tensor,
    ) -> torch.Tensor:
        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)  # (B, S, h, N, d)
        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)

        scale = 1.0 / math.sqrt(d)
        
        # Call fused attention kernel
        out = fused_ops.fused_attention_cuda(
            q.contiguous(), k.contiguous(), v.contiguous(), g.contiguous(),
            bias.contiguous(), msa_mask.bool(), B, S, N, h, d, scale
        )
        
        # out is (B, S, h, N, d), need to permute to (B, S, N, h, d) then reshape
        out = out.permute(0, 1, 3, 2, 4).reshape(B, S, N, D)
        return self.out_proj(out)