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


class Model(nn.Module):
    """
    Contact map prediction head.  Takes residue-pair embeddings and
    predicts a symmetric binary contact probability matrix (residue i
    within 8 Å of residue j).  Uses a small MLP on the symmetrised pair
    representation followed by a sigmoid.  Common final stage in protein
    structure prediction pipelines.
    """

    def __init__(self, pair_dim, hidden_dim):
        super().__init__()
        self.layer_norm = nn.LayerNorm(pair_dim)
        self.mlp = nn.Sequential(
            nn.Linear(pair_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1),
        )

    def forward(self, pair_repr: torch.Tensor) -> torch.Tensor:
        """
        Args:
            pair_repr: (B, N, N, pair_dim)
        Returns:
            contact_probs: (B, N, N) – symmetric contact probability
        """
        x = self.layer_norm(pair_repr)
        # Symmetrise
        x = (x + x.permute(0, 2, 1, 3)) / 2.0
        logits = self.mlp(x).squeeze(-1)  # (B, N, N)
        return torch.sigmoid(logits)


pair_dim = 64
hidden_dim = 32
seq_len = 64
batch_size = 4


def get_inputs():
    return [torch.randn(popcorn_pri.jitter_int(batch_size), popcorn_pri.jitter_int(seq_len), popcorn_pri.jitter_int(seq_len), pair_dim)]


def get_init_inputs():
    return [pair_dim, hidden_dim]

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

cuda_source = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>

// Fused kernel: LayerNorm + Symmetrize + MLP (3 layers) + Sigmoid
template <typename scalar_t>
__global__ void fused_contact_map_kernel(
    const scalar_t* __restrict__ pair_repr,
    const scalar_t* __restrict__ ln_weight,
    const scalar_t* __restrict__ ln_bias,
    const scalar_t* __restrict__ fc1_weight,
    const scalar_t* __restrict__ fc1_bias,
    const scalar_t* __restrict__ fc2_weight,
    const scalar_t* __restrict__ fc2_bias,
    const scalar_t* __restrict__ fc3_weight,
    const scalar_t* __restrict__ fc3_bias,
    scalar_t* __restrict__ output,
    const int batch_size,
    const int seq_len,
    const int pair_dim,
    const int hidden_dim) {
    
    // Each thread processes one output element (b, i, j)
    const int b = blockIdx.x;
    const int ij = blockIdx.y * blockDim.x + threadIdx.x;
    const int i = ij / seq_len;
    const int j = ij % seq_len;
    
    if (b >= batch_size || i >= seq_len || j >= seq_len) return;
    
    // Index calculation for pair_repr[b,i,j,:] and pair_repr[b,j,i,:]
    const int stride_seq = seq_len * pair_dim;
    const int stride_pair = pair_dim;
    const int base_b = b * seq_len * seq_len * pair_dim;
    const int base_idx_ij = base_b + i * stride_seq + j * stride_pair;
    const int base_idx_ji = base_b + j * stride_seq + i * stride_pair;
    
    // Compute LayerNorm for position (i,j)
    scalar_t mean_ij = 0;
    for (int d = 0; d < pair_dim; d++) {
        mean_ij += pair_repr[base_idx_ij + d];
    }
    mean_ij /= scalar_t(pair_dim);
    
    scalar_t var_ij = 0;
    for (int d = 0; d < pair_dim; d++) {
        scalar_t diff = pair_repr[base_idx_ij + d] - mean_ij;
        var_ij += diff * diff;
    }
    var_ij /= scalar_t(pair_dim);
    
    scalar_t inv_std_ij = rsqrt(var_ij + scalar_t(1e-5));
    
    // Compute LayerNorm for position (j,i)
    scalar_t mean_ji = 0;
    for (int d = 0; d < pair_dim; d++) {
        mean_ji += pair_repr[base_idx_ji + d];
    }
    mean_ji /= scalar_t(pair_dim);
    
    scalar_t var_ji = 0;
    for (int d = 0; d < pair_dim; d++) {
        scalar_t diff = pair_repr[base_idx_ji + d] - mean_ji;
        var_ji += diff * diff;
    }
    var_ji /= scalar_t(pair_dim);
    
    scalar_t inv_std_ji = rsqrt(var_ji + scalar_t(1e-5));
    
    // Symmetrize: average of LayerNorm outputs
    // Store symmetrized values in registers
    scalar_t sym_pair[64];  // Assuming pair_dim <= 64
    for (int d = 0; d < pair_dim; d++) {
        scalar_t ln_ij = (pair_repr[base_idx_ij + d] - mean_ij) * inv_std_ij * ln_weight[d] + ln_bias[d];
        scalar_t ln_ji = (pair_repr[base_idx_ji + d] - mean_ji) * inv_std_ji * ln_weight[d] + ln_bias[d];
        sym_pair[d] = (ln_ij + ln_ji) * scalar_t(0.5);
    }
    
    // MLP Layer 1: Linear(pair_dim -> hidden_dim) + ReLU
    scalar_t hidden1[32];  // Assuming hidden_dim <= 32
    for (int h = 0; h < hidden_dim; h++) {
        scalar_t sum = fc1_bias[h];
        for (int d = 0; d < pair_dim; d++) {
            sum += sym_pair[d] * fc1_weight[h * pair_dim + d];
        }
        hidden1[h] = max(scalar_t(0), sum);
    }
    
    // MLP Layer 2: Linear(hidden_dim -> hidden_dim) + ReLU
    scalar_t hidden2[32];
    for (int h = 0; h < hidden_dim; h++) {
        scalar_t sum = fc2_bias[h];
        for (int p = 0; p < hidden_dim; p++) {
            sum += hidden1[p] * fc2_weight[h * hidden_dim + p];
        }
        hidden2[h] = max(scalar_t(0), sum);
    }
    
    // MLP Layer 3: Linear(hidden_dim -> 1)
    scalar_t logit = fc3_bias[0];
    for (int h = 0; h < hidden_dim; h++) {
        logit += hidden2[h] * fc3_weight[h];
    }
    
    // Sigmoid
    output[b * seq_len * seq_len + i * seq_len + j] = scalar_t(1.0) / (scalar_t(1.0) + exp(-logit));
}

torch::Tensor fused_contact_map_cuda(
    torch::Tensor pair_repr,
    torch::Tensor ln_weight,
    torch::Tensor ln_bias,
    torch::Tensor fc1_weight,
    torch::Tensor fc1_bias,
    torch::Tensor fc2_weight,
    torch::Tensor fc2_bias,
    torch::Tensor fc3_weight,
    torch::Tensor fc3_bias) {
    
    const int batch_size = pair_repr.size(0);
    const int seq_len = pair_repr.size(1);
    const int pair_dim = pair_repr.size(3);
    const int hidden_dim = fc1_bias.size(0);
    
    auto output = torch::zeros({batch_size, seq_len, seq_len}, pair_repr.options());
    
    const int threads_per_block = 256;
    const dim3 blocks(batch_size, (seq_len * seq_len + threads_per_block - 1) / threads_per_block);
    
    AT_DISPATCH_FLOATING_TYPES(pair_repr.scalar_type(), "fused_contact_map_kernel", ([&] {
        fused_contact_map_kernel<scalar_t><<<blocks, threads_per_block>>>(
            pair_repr.data_ptr<scalar_t>(),
            ln_weight.data_ptr<scalar_t>(),
            ln_bias.data_ptr<scalar_t>(),
            fc1_weight.data_ptr<scalar_t>(),
            fc1_bias.data_ptr<scalar_t>(),
            fc2_weight.data_ptr<scalar_t>(),
            fc2_bias.data_ptr<scalar_t>(),
            fc3_weight.data_ptr<scalar_t>(),
            fc3_bias.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            batch_size, seq_len, pair_dim, hidden_dim);
    }));
    
    return output;
}
"""

cpp_source = "torch::Tensor fused_contact_map_cuda(torch::Tensor pair_repr, torch::Tensor ln_weight, torch::Tensor ln_bias, torch::Tensor fc1_weight, torch::Tensor fc1_bias, torch::Tensor fc2_weight, torch::Tensor fc2_bias, torch::Tensor fc3_weight, torch::Tensor fc3_bias);"

fused_module = load_inline(
    name='fused_contact_map_v2',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['fused_contact_map_cuda'],
    extra_cuda_cflags=['-O3'],
    verbose=True
)

class ModelNew(nn.Module):
    def __init__(self, pair_dim, hidden_dim):
        super().__init__()
        self.layer_norm = nn.LayerNorm(pair_dim)
        self.mlp = nn.Sequential(
            nn.Linear(pair_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1),
        )
        self.pair_dim = pair_dim
        self.hidden_dim = hidden_dim

    def forward(self, pair_repr: torch.Tensor) -> torch.Tensor:
        pair_repr = pair_repr.contiguous()
        
        ln_weight = self.layer_norm.weight.contiguous()
        ln_bias = self.layer_norm.bias.contiguous()
        
        fc1_weight = self.mlp[0].weight.contiguous()
        fc1_bias = self.mlp[0].bias.contiguous()
        fc2_weight = self.mlp[2].weight.contiguous()
        fc2_bias = self.mlp[2].bias.contiguous()
        fc3_weight = self.mlp[4].weight.contiguous()
        fc3_bias = self.mlp[4].bias.contiguous()
        
        return fused_module.fused_contact_map_cuda(
            pair_repr, ln_weight, ln_bias,
            fc1_weight, fc1_bias,
            fc2_weight, fc2_bias,
            fc3_weight, fc3_bias
        )