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


class MSARowSelfAttention(nn.Module):
    def __init__(self, msa_dim, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = msa_dim // num_heads
        self.norm = nn.LayerNorm(msa_dim)
        self.q = nn.Linear(msa_dim, msa_dim, bias=False)
        self.k = nn.Linear(msa_dim, msa_dim, bias=False)
        self.v = nn.Linear(msa_dim, msa_dim, bias=False)
        self.out = nn.Linear(msa_dim, msa_dim)

    def forward(self, msa):
        B, S, N, D = msa.shape
        h, d = self.num_heads, self.head_dim
        x = self.norm(msa)
        q = self.q(x).view(B, S, N, h, d).permute(0, 1, 3, 2, 4)
        k = self.k(x).view(B, S, N, h, d).permute(0, 1, 3, 2, 4)
        v = self.v(x).view(B, S, N, h, d).permute(0, 1, 3, 2, 4)
        attn = (q @ k.transpose(-1, -2)) / math.sqrt(d)
        attn = F.softmax(attn, dim=-1)
        out = (attn @ v).permute(0, 1, 3, 2, 4).reshape(B, S, N, D)
        return self.out(out)


class OuterProductMeanSmall(nn.Module):
    def __init__(self, msa_dim, pair_dim, hidden=32):
        super().__init__()
        self.norm = nn.LayerNorm(msa_dim)
        self.a = nn.Linear(msa_dim, hidden)
        self.b = nn.Linear(msa_dim, hidden)
        self.out = nn.Linear(hidden * hidden, pair_dim)

    def forward(self, msa):
        x = self.norm(msa)
        a = self.a(x)
        b = self.b(x)
        outer = torch.einsum("bsih,bsjk->bijhk", a, b)
        S = msa.shape[1]
        outer = outer / max(S, 1)
        B, N1, N2, H1, H2 = outer.shape
        return self.out(outer.reshape(B, N1, N2, H1 * H2))


class TriMulOut(nn.Module):
    def __init__(self, pair_dim, hidden):
        super().__init__()
        self.norm = nn.LayerNorm(pair_dim)
        self.pa = nn.Linear(pair_dim, hidden)
        self.pb = nn.Linear(pair_dim, hidden)
        self.ga = nn.Linear(pair_dim, hidden)
        self.gb = nn.Linear(pair_dim, hidden)
        self.norm_out = nn.LayerNorm(hidden)
        self.out = nn.Linear(hidden, pair_dim)
        self.gate_out = nn.Linear(pair_dim, pair_dim)

    def forward(self, pair):
        x = self.norm(pair)
        a = self.pa(x) * torch.sigmoid(self.ga(x))
        b = self.pb(x) * torch.sigmoid(self.gb(x))
        o = torch.einsum("bikh,bjkh->bijh", a, b)
        o = self.out(self.norm_out(o))
        return torch.sigmoid(self.gate_out(pair)) * o


class Model(nn.Module):
    """
    Full Evoformer block from AlphaFold2 combining MSA row self-attention,
    outer-product mean update, and triangular multiplicative update of the
    pair representation.  A single block of the iterative refinement stack.
    """

    def __init__(self, msa_dim, pair_dim, num_heads, tri_hidden):
        super().__init__()
        self.msa_attn = MSARowSelfAttention(msa_dim, num_heads)
        self.opm = OuterProductMeanSmall(msa_dim, pair_dim)
        self.tri_mul = TriMulOut(pair_dim, tri_hidden)

    def forward(
        self, msa_repr: torch.Tensor, pair_repr: torch.Tensor
    ) -> tuple:
        """
        Args:
            msa_repr:  (B, S, N, msa_dim)
            pair_repr: (B, N, N, pair_dim)
        Returns:
            (updated_msa, updated_pair)
        """
        msa_repr = msa_repr + self.msa_attn(msa_repr)
        pair_repr = pair_repr + self.opm(msa_repr)
        pair_repr = pair_repr + self.tri_mul(pair_repr)
        return msa_repr, pair_repr


msa_dim = 64
pair_dim = 64
num_heads = 4
tri_hidden = 32
seq_len = 32
msa_depth = 8
batch_size = 2


def get_inputs():
    p = popcorn_pri
    bs = p.jitter_int(batch_size)
    sd = p.jitter_int(msa_depth)
    sl = p.jitter_int(seq_len)
    msa = torch.randn(bs, sd, sl, msa_dim)
    pair = torch.randn(bs, sl, sl, pair_dim)
    return [msa, pair]


def get_init_inputs():
    return [msa_dim, pair_dim, num_heads, tri_hidden]

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from torch.utils.cpp_extension import load_inline

# Custom CUDA kernels for fused operations
cuda_code = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

#define WARP_SIZE 32
#define CUDA_CHECK(expr) do { cudaError_t err = (expr); if (err != cudaSuccess) { /* handle error */ } } while(0)

// Warp-level reduction for sum
__device__ __forceinline__ float warp_reduce_sum(float val) {
    #pragma unroll
    for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) {
        val += __shfl_xor_sync(0xffffffff, val, mask);
    }
    return val;
}

// Block-level reduction using shared memory
__device__ __forceinline__ float block_reduce_sum(float val, float* shared) {
    int lane = threadIdx.x % WARP_SIZE;
    int wid = threadIdx.x / WARP_SIZE;
    
    val = warp_reduce_sum(val);
    
    if (lane == 0) shared[wid] = val;
    __syncthreads();
    
    if (wid == 0) {
        val = (lane < (blockDim.x + WARP_SIZE - 1) / WARP_SIZE) ? shared[lane] : 0.0f;
        val = warp_reduce_sum(val);
    }
    __syncthreads();
    
    return val;
}

// Fused LayerNorm + 3 Linear projections (for Q, K, V)
__global__ void layernorm_qkv_kernel(
    const float* __restrict__ input,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ q_weight,
    const float* __restrict__ k_weight,
    const float* __restrict__ v_weight,
    float* __restrict__ q_out,
    float* __restrict__ k_out,
    float* __restrict__ v_out,
    int total_positions, int D_in, int D_out, float eps) {
    
    extern __shared__ float shared_mem[];
    float* shared = shared_mem;
    
    int pos = blockIdx.x;
    if (pos >= total_positions) return;
    
    const float* in_ptr = input + pos * D_in;
    float* q_ptr = q_out + pos * D_out;
    float* k_ptr = k_out + pos * D_out;
    float* v_ptr = v_out + pos * D_out;
    
    // Compute mean
    float sum = 0.0f;
    for (int i = threadIdx.x; i < D_in; i += blockDim.x) {
        sum += in_ptr[i];
    }
    sum = block_reduce_sum(sum, shared);
    float mean = sum / D_in;
    
    // Compute variance
    float sq_sum = 0.0f;
    for (int i = threadIdx.x; i < D_in; i += blockDim.x) {
        float diff = in_ptr[i] - mean;
        sq_sum += diff * diff;
    }
    sq_sum = block_reduce_sum(sq_sum, shared);
    float inv_std = rsqrtf(sq_sum / D_in + eps);
    
    // Compute normalized values and store in registers (for small D_in)
    // Then compute Q, K, V projections
    for (int o = threadIdx.x; o < D_out; o += blockDim.x) {
        float q_val = 0.0f, k_val = 0.0f, v_val = 0.0f;
        for (int i = 0; i < D_in; i++) {
            float norm_val = (in_ptr[i] - mean) * inv_std * gamma[i] + beta[i];
            q_val += norm_val * q_weight[o * D_in + i];
            k_val += norm_val * k_weight[o * D_in + i];
            v_val += norm_val * v_weight[o * D_in + i];
        }
        q_ptr[o] = q_val;
        k_ptr[o] = k_val;
        v_ptr[o] = v_val;
    }
}

// Fused LayerNorm + 2 Linear projections (for OPM a and b)
__global__ void layernorm_ab_kernel(
    const float* __restrict__ input,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ a_weight,
    const float* __restrict__ b_weight,
    float* __restrict__ a_out,
    float* __restrict__ b_out,
    int total_positions, int D_in, int D_out, float eps) {
    
    extern __shared__ float shared_mem[];
    float* shared = shared_mem;
    
    int pos = blockIdx.x;
    if (pos >= total_positions) return;
    
    const float* in_ptr = input + pos * D_in;
    float* a_ptr = a_out + pos * D_out;
    float* b_ptr = b_out + pos * D_out;
    
    // Compute mean
    float sum = 0.0f;
    for (int i = threadIdx.x; i < D_in; i += blockDim.x) {
        sum += in_ptr[i];
    }
    sum = block_reduce_sum(sum, shared);
    float mean = sum / D_in;
    
    // Compute variance
    float sq_sum = 0.0f;
    for (int i = threadIdx.x; i < D_in; i += blockDim.x) {
        float diff = in_ptr[i] - mean;
        sq_sum += diff * diff;
    }
    sq_sum = block_reduce_sum(sq_sum, shared);
    float inv_std = rsqrtf(sq_sum / D_in + eps);
    
    // Compute normalized values and projections
    for (int o = threadIdx.x; o < D_out; o += blockDim.x) {
        float a_val = 0.0f, b_val = 0.0f;
        for (int i = 0; i < D_in; i++) {
            float norm_val = (in_ptr[i] - mean) * inv_std * gamma[i] + beta[i];
            a_val += norm_val * a_weight[o * D_in + i];
            b_val += norm_val * b_weight[o * D_in + i];
        }
        a_ptr[o] = a_val;
        b_ptr[o] = b_val;
    }
}

// Fused LayerNorm + 4 Linear projections (for TriMul pa, pb, ga, gb)
__global__ void layernorm_tri_kernel(
    const float* __restrict__ input,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    const float* __restrict__ pa_weight,
    const float* __restrict__ pb_weight,
    const float* __restrict__ ga_weight,
    const float* __restrict__ gb_weight,
    float* __restrict__ a_out,
    float* __restrict__ b_out,
    int total_positions, int D_in, int D_out, float eps) {
    
    extern __shared__ float shared_mem[];
    float* shared = shared_mem;
    
    int pos = blockIdx.x;
    if (pos >= total_positions) return;
    
    const float* in_ptr = input + pos * D_in;
    float* a_ptr = a_out + pos * D_out;
    float* b_ptr = b_out + pos * D_out;
    
    // Compute mean
    float sum = 0.0f;
    for (int i = threadIdx.x; i < D_in; i += blockDim.x) {
        sum += in_ptr[i];
    }
    sum = block_reduce_sum(sum, shared);
    float mean = sum / D_in;
    
    // Compute variance
    float sq_sum = 0.0f;
    for (int i = threadIdx.x; i < D_in; i += blockDim.x) {
        float diff = in_ptr[i] - mean;
        sq_sum += diff * diff;
    }
    sq_sum = block_reduce_sum(sq_sum, shared);
    float inv_std = rsqrtf(sq_sum / D_in + eps);
    
    // Compute normalized values and projections with gating
    for (int o = threadIdx.x; o < D_out; o += blockDim.x) {
        float pa_val = 0.0f, pb_val = 0.0f;
        float ga_val = 0.0f, gb_val = 0.0f;
        for (int i = 0; i < D_in; i++) {
            float norm_val = (in_ptr[i] - mean) * inv_std * gamma[i] + beta[i];
            pa_val += norm_val * pa_weight[o * D_in + i];
            pb_val += norm_val * pb_weight[o * D_in + i];
            ga_val += norm_val * ga_weight[o * D_in + i];
            gb_val += norm_val * gb_weight[o * D_in + i];
        }
        // Apply sigmoid gating: a = pa * sigmoid(ga), b = pb * sigmoid(gb)
        a_ptr[o] = pa_val * (1.0f / (1.0f + expf(-ga_val)));
        b_ptr[o] = pb_val * (1.0f / (1.0f + expf(-gb_val)));
    }
}

// C++ interface functions
torch::Tensor layernorm_qkv_cuda(
    torch::Tensor input, torch::Tensor gamma, torch::Tensor beta,
    torch::Tensor q_weight, torch::Tensor k_weight, torch::Tensor v_weight,
    int D_out, float eps) {
    
    int total_positions = input.numel() / input.size(-1);
    int D_in = input.size(-1);
    
    auto q_out = torch::empty({total_positions, D_out}, input.options());
    auto k_out = torch::empty({total_positions, D_out}, input.options());
    auto v_out = torch::empty({total_positions, D_out}, input.options());
    
    int threads = (D_in + 31) / 32 * 32;
    threads = min(threads, 256);
    int blocks = total_positions;
    int shared_mem = ((threads + 31) / 32) * sizeof(float);
    
    layernorm_qkv_kernel<<<blocks, threads, shared_mem>>>(
        input.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(),
        q_weight.data_ptr<float>(), k_weight.data_ptr<float>(), v_weight.data_ptr<float>(),
        q_out.data_ptr<float>(), k_out.data_ptr<float>(), v_out.data_ptr<float>(),
        total_positions, D_in, D_out, eps);
    
    return torch::stack({q_out, k_out, v_out}, 0);
}

torch::Tensor layernorm_ab_cuda(
    torch::Tensor input, torch::Tensor gamma, torch::Tensor beta,
    torch::Tensor a_weight, torch::Tensor b_weight,
    int D_out, float eps) {
    
    int total_positions = input.numel() / input.size(-1);
    int D_in = input.size(-1);
    
    auto a_out = torch::empty({total_positions, D_out}, input.options());
    auto b_out = torch::empty({total_positions, D_out}, input.options());
    
    int threads = (D_in + 31) / 32 * 32;
    threads = min(threads, 256);
    int blocks = total_positions;
    int shared_mem = ((threads + 31) / 32) * sizeof(float);
    
    layernorm_ab_kernel<<<blocks, threads, shared_mem>>>(
        input.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(),
        a_weight.data_ptr<float>(), b_weight.data_ptr<float>(),
        a_out.data_ptr<float>(), b_out.data_ptr<float>(),
        total_positions, D_in, D_out, eps);
    
    return torch::stack({a_out, b_out}, 0);
}

torch::Tensor layernorm_tri_cuda(
    torch::Tensor input, torch::Tensor gamma, torch::Tensor beta,
    torch::Tensor pa_weight, torch::Tensor pb_weight,
    torch::Tensor ga_weight, torch::Tensor gb_weight,
    int D_out, float eps) {
    
    int total_positions = input.numel() / input.size(-1);
    int D_in = input.size(-1);
    
    auto a_out = torch::empty({total_positions, D_out}, input.options());
    auto b_out = torch::empty({total_positions, D_out}, input.options());
    
    int threads = (D_in + 31) / 32 * 32;
    threads = min(threads, 256);
    int blocks = total_positions;
    int shared_mem = ((threads + 31) / 32) * sizeof(float);
    
    layernorm_tri_kernel<<<blocks, threads, shared_mem>>>(
        input.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(),
        pa_weight.data_ptr<float>(), pb_weight.data_ptr<float>(),
        ga_weight.data_ptr<float>(), gb_weight.data_ptr<float>(),
        a_out.data_ptr<float>(), b_out.data_ptr<float>(),
        total_positions, D_in, D_out, eps);
    
    return torch::stack({a_out, b_out}, 0);
}
"""

# Load the CUDA extension
fused_ops = load_inline(
    name='fused_ops_v1',
    cpp_sources='''
#include <torch/extension.h>
torch::Tensor layernorm_qkv_cuda(torch::Tensor input, torch::Tensor gamma, torch::Tensor beta,
    torch::Tensor q_weight, torch::Tensor k_weight, torch::Tensor v_weight, int D_out, float eps);
torch::Tensor layernorm_ab_cuda(torch::Tensor input, torch::Tensor gamma, torch::Tensor beta,
    torch::Tensor a_weight, torch::Tensor b_weight, int D_out, float eps);
torch::Tensor layernorm_tri_cuda(torch::Tensor input, torch::Tensor gamma, torch::Tensor beta,
    torch::Tensor pa_weight, torch::Tensor pb_weight, torch::Tensor ga_weight, torch::Tensor gb_weight,
    int D_out, float eps);
''',
    cuda_sources=cuda_code,
    functions=['layernorm_qkv_cuda', 'layernorm_ab_cuda', 'layernorm_tri_cuda'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    verbose=False
)


class ModelNew(nn.Module):
    """
    Optimized Evoformer block with custom CUDA kernels for fused operations.
    """
    
    def __init__(self, msa_dim, pair_dim, num_heads, tri_hidden):
        super().__init__()
        self.msa_dim = msa_dim
        self.pair_dim = pair_dim
        self.num_heads = num_heads
        self.head_dim = msa_dim // num_heads
        self.tri_hidden = tri_hidden
        self.eps = 1e-5
        
        # MSA attention weights (LayerNorm + Q/K/V projections)
        self.msa_ln_weight = nn.Parameter(torch.ones(msa_dim))
        self.msa_ln_bias = nn.Parameter(torch.zeros(msa_dim))
        self.q_weight = nn.Parameter(torch.randn(msa_dim, msa_dim) * (1.0 / math.sqrt(msa_dim)))
        self.k_weight = nn.Parameter(torch.randn(msa_dim, msa_dim) * (1.0 / math.sqrt(msa_dim)))
        self.v_weight = nn.Parameter(torch.randn(msa_dim, msa_dim) * (1.0 / math.sqrt(msa_dim)))
        self.out_weight = nn.Parameter(torch.randn(msa_dim, msa_dim) * (1.0 / math.sqrt(msa_dim)))
        self.out_bias = nn.Parameter(torch.zeros(msa_dim))
        
        # OPM weights (LayerNorm + a/b projections)
        self.opm_ln_weight = nn.Parameter(torch.ones(msa_dim))
        self.opm_ln_bias = nn.Parameter(torch.zeros(msa_dim))
        self.a_weight = nn.Parameter(torch.randn(msa_dim, 32) * (1.0 / math.sqrt(msa_dim)))
        self.b_weight = nn.Parameter(torch.randn(msa_dim, 32) * (1.0 / math.sqrt(msa_dim)))
        self.opm_out_weight = nn.Parameter(torch.randn(32 * 32, pair_dim) * (1.0 / math.sqrt(32 * 32)))
        self.opm_out_bias = nn.Parameter(torch.zeros(pair_dim))
        
        # TriMul weights (LayerNorm + pa/pb/ga/gb projections)
        self.tri_ln_weight = nn.Parameter(torch.ones(pair_dim))
        self.tri_ln_bias = nn.Parameter(torch.zeros(pair_dim))
        self.pa_weight = nn.Parameter(torch.randn(pair_dim, tri_hidden) * (1.0 / math.sqrt(pair_dim)))
        self.pb_weight = nn.Parameter(torch.randn(pair_dim, tri_hidden) * (1.0 / math.sqrt(pair_dim)))
        self.ga_weight = nn.Parameter(torch.randn(pair_dim, tri_hidden) * (1.0 / math.sqrt(pair_dim)))
        self.gb_weight = nn.Parameter(torch.randn(pair_dim, tri_hidden) * (1.0 / math.sqrt(pair_dim)))
        self.tri_out_ln_weight = nn.Parameter(torch.ones(tri_hidden))
        self.tri_out_ln_bias = nn.Parameter(torch.zeros(tri_hidden))
        self.tri_out_weight = nn.Parameter(torch.randn(tri_hidden, pair_dim) * (1.0 / math.sqrt(tri_hidden)))
        self.tri_out_bias = nn.Parameter(torch.zeros(pair_dim))
        self.gate_weight = nn.Parameter(torch.randn(pair_dim, pair_dim) * (1.0 / math.sqrt(pair_dim)))
        self.gate_bias = nn.Parameter(torch.zeros(pair_dim))
    
    def _msa_attention(self, msa):
        """MSA row self-attention with fused LayerNorm + Q/K/V kernel."""
        B, S, N, D = msa.shape
        h, d = self.num_heads, self.head_dim
        
        # Reshape for fused kernel: (B*S*N, D)
        x = msa.reshape(-1, D)
        
        # Fused LayerNorm + Q/K/V projection
        qkv = fused_ops.layernorm_qkv_cuda(
            x, self.msa_ln_weight, self.msa_ln_bias,
            self.q_weight, self.k_weight, self.v_weight, D, self.eps
        )
        q = qkv[0].view(B, S, N, h, d).permute(0, 1, 3, 2, 4)
        k = qkv[1].view(B, S, N, h, d).permute(0, 1, 3, 2, 4)
        v = qkv[2].view(B, S, N, h, d).permute(0, 1, 3, 2, 4)
        
        # Scaled dot-product attention
        scale = 1.0 / math.sqrt(d)
        attn = (q @ k.transpose(-1, -2)) * scale
        attn = F.softmax(attn, dim=-1)
        out = (attn @ v).permute(0, 1, 3, 2, 4).reshape(B, S, N, D)
        
        # Output projection
        return F.linear(out, self.out_weight, self.out_bias)
    
    def _outer_product_mean(self, msa):
        """Outer product mean with fused LayerNorm + a/b kernel."""
        B, S, N, D = msa.shape
        H = 32  # hidden dimension
        
        # Reshape for fused kernel: (B*S*N, D)
        x = msa.reshape(-1, D)
        
        # Fused LayerNorm + a/b projection
        ab = fused_ops.layernorm_ab_cuda(
            x, self.opm_ln_weight, self.opm_ln_bias,
            self.a_weight, self.b_weight, H, self.eps
        )
        a = ab[0].view(B, S, N, H)
        b = ab[1].view(B, S, N, H)
        
        # Outer product: einsum("bsih,bsjk->bijhk", a, b)
        # Optimized using reshape + matmul
        outer = torch.einsum("bsih,bsjk->bijhk", a, b)
        outer = outer / max(S, 1)
        
        # Output projection
        B_out, N1, N2, H1, H2 = outer.shape
        outer_flat = outer.reshape(B_out, N1, N2, H1 * H2)
        return F.linear(outer_flat, self.opm_out_weight, self.opm_out_bias)
    
    def _tri_mul_out(self, pair):
        """Triangular multiplicative update with fused LayerNorm + gated projections."""
        B, N, _, D_pair = pair.shape
        H = self.tri_hidden
        
        # Reshape for fused kernel: (B*N*N, D_pair)
        x = pair.reshape(-1, D_pair)
        
        # Fused LayerNorm + gated pa/pb projections
        ab = fused_ops.layernorm_tri_cuda(
            x, self.tri_ln_weight, self.tri_ln_bias,
            self.pa_weight, self.pb_weight,
            self.ga_weight, self.gb_weight, H, self.eps
        )
        a = ab[0].view(B, N, N, H)
        b = ab[1].view(B, N, N, H)
        
        # Triangular multiplication: einsum("bikh,bjkh->bijh", a, b)
        o = torch.einsum("bikh,bjkh->bijh", a, b)
        
        # LayerNorm + output projection
        o = F.layer_norm(o, (H,), self.tri_out_ln_weight, self.tri_out_ln_bias)
        o = F.linear(o, self.tri_out_weight, self.tri_out_bias)
        
        # Gated residual connection
        gate = torch.sigmoid(F.linear(pair, self.gate_weight, self.gate_bias))
        return gate * o
    
    def forward(self, msa_repr: torch.Tensor, pair_repr: torch.Tensor) -> tuple:
        """
        Forward pass with fused CUDA kernels.
        
        Args:
            msa_repr:  (B, S, N, msa_dim)
            pair_repr: (B, N, N, pair_dim)
        Returns:
            (updated_msa, updated_pair)
        """
        # MSA row self-attention
        msa_repr = msa_repr + self._msa_attention(msa_repr)
        
        # Outer product mean
        pair_repr = pair_repr + self._outer_product_mean(msa_repr)
        
        # Triangular multiplicative update
        pair_repr = pair_repr + self._tri_mul_out(pair_repr)
        
        return msa_repr, pair_repr