import torch
import torch.nn as nn
import torch.nn.functional as F
import math


class Model(nn.Module):
    """
    Triangular self-attention over pair representations, as used in
    AlphaFold2's Evoformer stack.  Starting-node variant: for each
    row i of the pair matrix, standard multi-head self-attention is
    performed across the column dimension j, with an additive pair
    bias.  This lets positions (i,j) and (i,k) exchange information
    through their shared starting node i.
    """

    def __init__(self, pair_dim, num_heads, dropout=0.0):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = pair_dim // num_heads
        assert pair_dim % num_heads == 0

        self.layer_norm = nn.LayerNorm(pair_dim)
        self.query = nn.Linear(pair_dim, pair_dim, bias=False)
        self.key = nn.Linear(pair_dim, pair_dim, bias=False)
        self.value = nn.Linear(pair_dim, pair_dim, bias=False)
        self.gate = nn.Linear(pair_dim, pair_dim)
        self.bias_proj = nn.Linear(pair_dim, num_heads, bias=False)
        self.out_proj = nn.Linear(pair_dim, pair_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, pair_repr: torch.Tensor, pair_mask: torch.Tensor) -> torch.Tensor:
        """
        Args:
            pair_repr: (B, N, N, pair_dim)
            pair_mask: (B, N, N)
        Returns:
            updated pair representation (B, N, N, pair_dim)
        """
        B, N, _, D = pair_repr.shape
        h = self.num_heads
        d = self.head_dim

        x = self.layer_norm(pair_repr)
        g = torch.sigmoid(self.gate(x))

        # Pair bias from the pair representation itself
        bias = self.bias_proj(x).permute(0, 3, 1, 2)  # (B, h, N, N)

        # Reshape to (B*N, N, D) — attention across columns for each row
        x_flat = x.reshape(B * N, N, D)
        q = self.query(x_flat).view(B * N, N, h, d).transpose(1, 2)  # (B*N, h, N, d)
        k = self.key(x_flat).view(B * N, N, h, d).transpose(1, 2)
        v = self.value(x_flat).view(B * N, N, h, d).transpose(1, 2)

        attn = (q @ k.transpose(-1, -2)) / math.sqrt(d)  # (B*N, h, N, N)
        attn = attn.view(B, N, h, N, N)

        # Add pair bias: bias (B, h, N, N) → (B, 1, h, N, N) broadcast across rows
        attn = attn + bias.unsqueeze(1)

        # Apply mask: pair_mask (B, N, N) → mask over the key dimension
        mask = pair_mask.unsqueeze(1).unsqueeze(2)  # (B, 1, 1, N, N) – last dim is keys
        attn = attn.masked_fill(mask == 0, float("-inf"))
        attn = attn.view(B * N, h, N, N)
        attn = F.softmax(attn, dim=-1)
        attn = self.dropout(attn)

        out = (attn @ v).transpose(1, 2).reshape(B, N, N, D)
        out = g * out
        return self.out_proj(out)


pair_dim = 64
num_heads = 4
seq_len = 32
batch_size = 2


def get_inputs():
    p = popcorn_pri
    pair_repr = torch.randn(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(seq_len), pair_dim)
    pair_mask = torch.ones(p.jitter_int(batch_size), p.jitter_int(seq_len), p.jitter_int(seq_len))
    return [pair_repr, pair_mask]


def get_init_inputs():
    return [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

# Define CUDA kernels
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

// Fused attention kernel - each block handles one (batch_row, head)
// Threads cooperate to process all query positions for this block
template <typename scalar_t>
__global__ void fused_attention_kernel(
    const scalar_t* __restrict__ query,    // (BN, h, N, d) - contiguous
    const scalar_t* __restrict__ key,      // (BN, h, N, d)
    const scalar_t* __restrict__ value,    // (BN, h, N, d)
    const scalar_t* __restrict__ bias,     // (B, h, N, N)
    const bool* __restrict__ mask,         // (B, N, N)
    scalar_t* __restrict__ output,         // (BN, N, h, d)
    const int B, const int N, const int h, const int d, const int BN,
    const float scale
) {
    // Block index determines which (batch_row, head) we handle
    const int batch_row_idx = blockIdx.x / h;
    const int head_idx = blockIdx.x % h;
    const int batch_idx = batch_row_idx / N;
    
    const int tid = threadIdx.x;
    const int blockDimX = blockDim.x;
    
    // Shared memory layout:
    // - K tile: [TILE_SIZE, d]
    // - V tile: [TILE_SIZE, d]
    // - attention row: [N] - one per query position, computed cooperatively
    extern __shared__ char smem[];
    scalar_t* k_tile = reinterpret_cast<scalar_t*>(smem);
    scalar_t* v_tile = k_tile + blockDimX * d;
    float* attn_row = reinterpret_cast<float*>(v_tile + blockDimX * d);
    
    // Base pointers for Q, K, V
    const scalar_t* q_base = query + (batch_row_idx * h + head_idx) * N * d;
    const scalar_t* k_base = key + (batch_row_idx * h + head_idx) * N * d;
    const scalar_t* v_base = value + (batch_row_idx * h + head_idx) * N * d;
    // Bias is indexed as bias[b, head, q_pos, k_pos]
    const scalar_t* bias_base = bias + batch_idx * h * N * N + head_idx * N * N;
    // Mask base for this batch
    const bool* mask_base = mask + batch_idx * N * N;
    
    // Process each query position sequentially (thread 0 drives, others help)
    for (int q_idx = 0; q_idx < N; q_idx++) {
        // Load query vector into shared memory (thread 0 does this)
        __shared__ float q_shared[64];  // d <= 64
        if (tid == 0) {
            const scalar_t* q_ptr = q_base + q_idx * d;
            for (int i = 0; i < d; i++) {
                q_shared[i] = float(q_ptr[i]);
            }
        }
        __syncthreads();
        
        // Thread-local accumulators
        float thread_max = -INFINITY;
        float thread_sum = 0.0f;
        
        // First pass: compute all attention scores and find max
        // Each thread computes a subset of the scores
        for (int k_pos = tid; k_pos < N; k_pos += blockDimX) {
            float score = 0.0f;
            // Compute dot product with key at position k_pos
            const scalar_t* k_ptr = k_base + k_pos * d;
            for (int i = 0; i < d; i++) {
                score += q_shared[i] * float(k_ptr[i]);
            }
            score *= scale;
            // Add bias
            score += float(bias_base[q_idx * N + k_pos]);
            // Apply mask
            if (!mask_base[q_idx * N + k_pos]) {
                score = -INFINITY;
            }
            attn_row[k_pos] = score;
            if (score > thread_max) {
                thread_max = score;
            }
        }
        __syncthreads();
        
        // Reduce max across threads (simple parallel reduction)
        __shared__ float shared_max;
        if (tid == 0) {
            shared_max = -INFINITY;
        }
        __syncthreads();
        
        // Each thread writes its max
        if (thread_max > -INFINITY/2) {
            atomicMax(reinterpret_cast<int*>(&shared_max), __float_as_int(thread_max));
        }
        __syncthreads();
        
        float max_score = shared_max;
        
        // Compute exp(scores - max) and sum
        for (int k_pos = tid; k_pos < N; k_pos += blockDimX) {
            float score = attn_row[k_pos];
            if (score > -INFINITY / 2) {
                float exp_val = expf(score - max_score);
                attn_row[k_pos] = exp_val;
                thread_sum += exp_val;
            } else {
                attn_row[k_pos] = 0.0f;
            }
        }
        
        // Reduce sum across threads
        __shared__ float shared_sum;
        if (tid == 0) {
            shared_sum = 0.0f;
        }
        __syncthreads();
        
        // Use atomic add for sum reduction
        if (thread_sum > 0) {
            atomicAdd(&shared_sum, thread_sum);
        }
        __syncthreads();
        
        float sum_exp = shared_sum;
        float inv_sum = 1.0f / (sum_exp + 1e-8f);
        
        // Normalize
        for (int k_pos = tid; k_pos < N; k_pos += blockDimX) {
            attn_row[k_pos] *= inv_sum;
        }
        __syncthreads();
        
        // Second pass: compute weighted sum of values
        // Each thread computes partial sum, then we reduce
        float thread_out[64];  // d <= 64
        for (int i = 0; i < d; i++) thread_out[i] = 0.0f;
        
        for (int k_pos = tid; k_pos < N; k_pos += blockDimX) {
            float attn_w = attn_row[k_pos];
            if (attn_w > 0) {
                const scalar_t* v_ptr = v_base + k_pos * d;
                for (int i = 0; i < d; i++) {
                    thread_out[i] += attn_w * float(v_ptr[i]);
                }
            }
        }
        __syncthreads();
        
        // Reduce output across threads
        __shared__ float shared_out[64];
        if (tid == 0) {
            for (int i = 0; i < d; i++) {
                shared_out[i] = 0.0f;
            }
        }
        __syncthreads();
        
        for (int i = 0; i < d; i++) {
            atomicAdd(&shared_out[i], thread_out[i]);
        }
        __syncthreads();
        
        // Thread 0 writes output
        if (tid == 0) {
            scalar_t* out_ptr = output + (batch_row_idx * N + q_idx) * h * d + head_idx * d;
            for (int i = 0; i < d; i++) {
                out_ptr[i] = scalar_t(shared_out[i]);
            }
        }
        __syncthreads();
    }
}

// Host function
torch::Tensor fused_attention_cuda(
    torch::Tensor query,
    torch::Tensor key,
    torch::Tensor value,
    torch::Tensor bias,
    torch::Tensor mask,
    int B, int N, int h, int d
) {
    const int BN = B * N;
    const float scale = 1.0f / sqrtf((float)d);
    
    // Output: (BN, N, h, d)
    auto output = torch::empty({BN, N, h, d}, query.options());
    
    const int num_blocks = BN * h;
    const int threads_per_block = 256;
    const int shared_mem_size = threads_per_block * d * 2 * sizeof(float) + N * sizeof(float);
    
    AT_DISPATCH_FLOATING_TYPES(query.scalar_type(), "fused_attention_kernel", ([&] {
        fused_attention_kernel<scalar_t><<<num_blocks, threads_per_block, shared_mem_size>>>(
            query.data_ptr<scalar_t>(),
            key.data_ptr<scalar_t>(),
            value.data_ptr<scalar_t>(),
            bias.data_ptr<scalar_t>(),
            mask.data_ptr<bool>(),
            output.data_ptr<scalar_t>(),
            B, N, h, d, BN, scale
        );
    }));
    
    return output;
}
"""

cpp_source = """
torch::Tensor fused_attention_cuda(
    torch::Tensor query,
    torch::Tensor key,
    torch::Tensor value,
    torch::Tensor bias,
    torch::Tensor mask,
    int B, int N, int h, int d);
"""

# Compile extension
fused_attention_module = load_inline(
    name='fused_attention_v4',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['fused_attention_cuda'],
    extra_cuda_cflags=['-O3'],
    verbose=True
)


class ModelNew(nn.Module):
    def __init__(self, pair_dim, num_heads, dropout=0.0):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = pair_dim // num_heads
        assert pair_dim % num_heads == 0

        self.layer_norm = nn.LayerNorm(pair_dim)
        self.query = nn.Linear(pair_dim, pair_dim, bias=False)
        self.key = nn.Linear(pair_dim, pair_dim, bias=False)
        self.value = nn.Linear(pair_dim, pair_dim, bias=False)
        self.gate = nn.Linear(pair_dim, pair_dim)
        self.bias_proj = nn.Linear(pair_dim, num_heads, bias=False)
        self.out_proj = nn.Linear(pair_dim, pair_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, pair_repr: torch.Tensor, pair_mask: torch.Tensor) -> torch.Tensor:
        B, N, _, D = pair_repr.shape
        h = self.num_heads
        d = self.head_dim

        x = self.layer_norm(pair_repr)
        g = torch.sigmoid(self.gate(x))

        # Pair bias from the pair representation itself
        bias = self.bias_proj(x).permute(0, 3, 1, 2)  # (B, h, N, N)

        # Reshape to (B*N, N, D) — attention across columns for each row
        x_flat = x.reshape(B * N, N, D)
        q = self.query(x_flat).view(B * N, N, h, d).transpose(1, 2)  # (B*N, h, N, d)
        k = self.key(x_flat).view(B * N, N, h, d).transpose(1, 2)
        v = self.value(x_flat).view(B * N, N, h, d).transpose(1, 2)

        # Use custom fused attention kernel
        q = q.contiguous()
        k = k.contiguous()
        v = v.contiguous()
        bias = bias.contiguous()
        mask_bool = (pair_mask == 1).contiguous()
        
        out = fused_attention_module.fused_attention_cuda(q, k, v, bias, mask_bool, B, N, h, d)
        out = out.reshape(B, N, N, D)
        
        out = g * out
        return self.out_proj(out)


pair_dim = 64
num_heads = 4
seq_len = 32
batch_size = 2


def get_inputs():
    pair_repr = torch.randn(batch_size, seq_len, seq_len, pair_dim)
    pair_mask = torch.ones(batch_size, seq_len, seq_len)
    return [pair_repr, pair_mask]


def get_init_inputs():
    return [pair_dim, num_heads]