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 = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cmath>

#define MAX_D 64
#define MAX_H 32

__global__ void contact_map_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;

    long long base = (long long)b * N * N * D;
    const float* vec = pair_repr + base + (i * N + j) * D;
    float mean = 0.0f;
    #pragma unroll
    for (int d = 0; d < D; d++) mean += vec[d];
    mean /= D;
    float var = 0.0f;
    #pragma unroll
    for (int d = 0; d < D; d++) {
        float diff = vec[d] - mean;
        var += diff * diff;
    }
    var /= D;
    float inv_std = rsqrtf(var + 1e-5f);

    float x[MAX_D];
    #pragma unroll
    for (int d = 0; d < D; d++) {
        x[d] = (vec[d] - mean) * inv_std * ln_weight[d] + ln_bias[d];
    }

    const float* vec2 = pair_repr + base + (j * N + i) * D;
    float mean2 = 0.0f;
    #pragma unroll
    for (int d = 0; d < D; d++) mean2 += vec2[d];
    mean2 /= D;
    float var2 = 0.0f;
    #pragma unroll
    for (int d = 0; d < D; d++) {
        float diff = vec2[d] - mean2;
        var2 += diff * diff;
    }
    var2 /= D;
    float inv_std2 = rsqrtf(var2 + 1e-5f);

    float x2[MAX_D];
    #pragma unroll
    for (int d = 0; d < D; d++) {
        x2[d] = (vec2[d] - mean2) * inv_std2 * ln_weight[d] + ln_bias[d];
    }

    float xs[MAX_D];
    #pragma unroll
    for (int d = 0; d < D; d++) {
        xs[d] = (x[d] + x2[d]) * 0.5f;
    }

    float h1[MAX_H];
    for (int h = 0; h < H; h++) {
        float sum = b1[h];
        #pragma unroll
        for (int d = 0; d < D; d++) {
            sum += xs[d] * w1[h * D + d];
        }
        h1[h] = fmaxf(sum, 0.0f);
    }

    float h2[MAX_H];
    for (int h = 0; h < H; h++) {
        float sum = b2[h];
        #pragma unroll
        for (int hh = 0; hh < H; hh++) {
            sum += h1[hh] * w2[h * H + hh];
        }
        h2[h] = fmaxf(sum, 0.0f);
    }

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

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

torch::Tensor contact_map_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;

    contact_map_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 contact_map_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="contact_map_cuda",
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=["contact_map_forward"],
    verbose=False,
    extra_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.cuda_module = module

    def forward(self, pair_repr: torch.Tensor) -> torch.Tensor:
        ln = self.layer_norm
        mlp = self.mlp
        return self.cuda_module.contact_map_forward(
            pair_repr,
            ln.weight,
            ln.bias,
            mlp[0].weight,
            mlp[0].bias,
            mlp[2].weight,
            mlp[2].bias,
            mlp[4].weight,
            mlp[4].bias
        )