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.h>
#include <cuda_runtime.h>
#include <cmath>

__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 pair_dim, int hidden_dim
) {
    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;

    const float* x = pair_repr + (b * N * N * pair_dim + (i * N + j) * pair_dim);
    const float* x_sym = pair_repr + (b * N * N * pair_dim + (j * N + i) * pair_dim);

    // LN on x and on x_sym separately, then average
    float mean1 = 0.0f, var1 = 0.0f;
    float mean2 = 0.0f, var2 = 0.0f;
    for (int d = 0; d < pair_dim; ++d) {
        mean1 += x[d];
        mean2 += x_sym[d];
    }
    mean1 /= pair_dim;
    mean2 /= pair_dim;
    for (int d = 0; d < pair_dim; ++d) {
        float dev1 = x[d] - mean1;
        float dev2 = x_sym[d] - mean2;
        var1 += dev1 * dev1;
        var2 += dev2 * dev2;
    }
    float inv_std1 = rsqrtf(var1 / pair_dim + 1e-5f);
    float inv_std2 = rsqrtf(var2 / pair_dim + 1e-5f);

    float norm1[64];
    float norm2[64];
    for (int d = 0; d < pair_dim; ++d) {
        norm1[d] = (x[d] - mean1) * inv_std1 * ln_weight[d] + ln_bias[d];
        norm2[d] = (x_sym[d] - mean2) * inv_std2 * ln_weight[d] + ln_bias[d];
    }

    // symmetrize
    float h1[32];
    for (int h = 0; h < hidden_dim; h++) {
        float sum = b1[h];
        for (int d = 0; d < pair_dim; d++) {
            float val = (norm1[d] + norm2[d]) * 0.5f;
            sum += val * w1[h * pair_dim + d];
        }
        h1[h] = fmaxf(sum, 0.0f);
    }

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

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

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

torch::Tensor contact_map_cuda(
    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 pair_dim = pair_repr.size(3);
    int hidden_dim = 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, pair_dim, hidden_dim
    );

    return output;
}
'''

cpp_source = '''
torch::Tensor contact_map_cuda(
    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',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['contact_map_cuda'],
    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.module = module

    def forward(self, pair_repr):
        ln = self.layer_norm
        mlp = self.mlp
        return self.module.contact_map_cuda(
            pair_repr,
            ln.weight, ln.bias,
            mlp[0].weight, mlp[0].bias,
            mlp[2].weight, mlp[2].bias,
            mlp[4].weight, mlp[4].bias
        )