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

cuda_source = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

__global__ void fused_contact_kernel(
    const float* __restrict__ pair_repr,
    const float* __restrict__ ln_weight,
    const float* __restrict__ ln_bias,
    const float* __restrict__ w1,
    const float* __restrict__ b1,
    const float* __restrict__ w2,
    const float* __restrict__ b2,
    const float* __restrict__ w3,
    const float* __restrict__ b3,
    float* __restrict__ output,
    int B, int N, int D, int H
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * N * N;
    if (idx >= total) return;

    int b = idx / (N * N);
    int rem = idx % (N * N);
    int i = rem / N;
    int j = rem % N;

    // LayerNorm for position (b,i,j) over D
    const float* x_ptr = pair_repr + (b * N * N + i * N + j) * D;
    float mean = 0.0f;
    for (int d = 0; d < D; d++) mean += x_ptr[d];
    mean /= D;

    float var = 0.0f;
    for (int d = 0; d < D; d++) {
        float diff = x_ptr[d] - mean;
        var += diff * diff;
    }
    var /= D;
    float inv_std = rsqrtf(var + 1e-5f);

    // Symmetrized LN: average with (j,i)
    const float* x_t_ptr = pair_repr + (b * N * N + j * N + i) * D;
    float mean_t = 0.0f;
    for (int d = 0; d < D; d++) mean_t += x_t_ptr[d];
    mean_t /= D;

    float var_t = 0.0f;
    for (int d = 0; d < D; d++) {
        float diff = x_t_ptr[d] - mean_t;
        var_t += diff * diff;
    }
    var_t /= D;
    float inv_std_t = rsqrtf(var_t + 1e-5f);

    // Compute averaged normalized features
    float norm_feat[64];
    for (int d = 0; d < D; d++) {
        float n1 = (x_ptr[d] - mean) * inv_std * ln_weight[d] + ln_bias[d];
        float n2 = (x_t_ptr[d] - mean_t) * inv_std_t * ln_weight[d] + ln_bias[d];
        norm_feat[d] = (n1 + n2) * 0.5f;
    }

    // MLP: D -> H -> H -> 1
    float h1[32];
    for (int h = 0; h < H; h++) {
        float sum = b1[h];
        for (int d = 0; d < D; d++) {
            sum += norm_feat[d] * w1[h * D + d];
        }
        h1[h] = fmaxf(sum, 0.0f);  // ReLU
    }

    float h2[32];
    for (int h = 0; h < H; h++) {
        float sum = b2[h];
        for (int hh = 0; hh < H; hh++) {
            sum += h1[hh] * w2[h * H + hh];
        }
        h2[h] = fmaxf(sum, 0.0f);  // ReLU
    }

    float logit = b3[0];
    for (int h = 0; h < H; h++) {
        logit += h2[h] * w3[h];
    }

    output[idx] = 1.0f / (1.0f + expf(-logit));  // sigmoid
}

torch::Tensor fused_contact_forward(
    torch::Tensor pair_repr,
    torch::Tensor ln_weight,
    torch::Tensor ln_bias,
    torch::Tensor w1, torch::Tensor b1,
    torch::Tensor w2, torch::Tensor b2,
    torch::Tensor w3, torch::Tensor b3
) {
    int B = pair_repr.size(0);
    int N = pair_repr.size(1);
    int D = pair_repr.size(3);
    int H = w1.size(0);

    auto output = torch::empty({B, N, N}, pair_repr.options());

    int total = B * N * N;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;

    fused_contact_kernel<<<blocks, threads>>>(
        pair_repr.data_ptr<float>(),
        ln_weight.data_ptr<float>(),
        ln_bias.data_ptr<float>(),
        w1.data_ptr<float>(),
        b1.data_ptr<float>(),
        w2.data_ptr<float>(),
        b2.data_ptr<float>(),
        w3.data_ptr<float>(),
        b3.data_ptr<float>(),
        output.data_ptr<float>(),
        B, N, D, H
    );

    return output;
}
'''

cpp_source = '''
torch::Tensor fused_contact_forward(
    torch::Tensor pair_repr,
    torch::Tensor ln_weight,
    torch::Tensor ln_bias,
    torch::Tensor w1, torch::Tensor b1,
    torch::Tensor w2, torch::Tensor b2,
    torch::Tensor w3, torch::Tensor b3
);
'''

module = load_inline(
    name='fused_contact',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['fused_contact_forward'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

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:
        ln = self.layer_norm
        mlp = self.mlp
        return module.fused_contact_forward(
            pair_repr.contiguous(),
            ln.weight, ln.bias,
            mlp[0].weight, mlp[0].bias,
            mlp[2].weight, mlp[2].bias,
            mlp[4].weight, mlp[4].bias
        )
'''