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

cpp_source = r"""
#include <torch/extension.h>
torch::Tensor 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);
"""

cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>
#define C_DIM 64
#define H_DIM 32
#define WARPS 4
#define BLOCK 128
__device__ __forceinline__ float warp_sum(float v){
    v += __shfl_down_sync(0xffffffff, v, 16);
    v += __shfl_down_sync(0xffffffff, v, 8);
    v += __shfl_down_sync(0xffffffff, v, 4);
    v += __shfl_down_sync(0xffffffff, v, 2);
    v += __shfl_down_sync(0xffffffff, v, 1);
    return v;
}
__global__ __launch_bounds__(BLOCK) void contact_kernel_tri(
    const float* __restrict__ inp, const float* __restrict__ ln_w, const float* __restrict__ ln_b,
    const float* __restrict__ w1, const float* __restrict__ bias1, const float* __restrict__ w2,
    const float* __restrict__ bias2, const float* __restrict__ w3, const float* __restrict__ bias3,
    float* __restrict__ out, int N, int tri, int total_tri)
{
    __shared__ float sx[WARPS*C_DIM];
    __shared__ float h1[WARPS*H_DIM];
    __shared__ float h2[WARPS*H_DIM];
    int lane = threadIdx.x & 31;
    int warp = threadIdx.x >> 5;
    int wid = blockIdx.x * WARPS + warp;
    if (wid >= total_tri) return;
    int b = wid / tri;
    int kk = wid - b * tri;
    float nf = (float)N;
    float val = (2.0f*nf + 1.0f);
    int i = (int)floorf((val - sqrtf(val*val - 8.0f*(float)kk)) * 0.5f);
    int start = i * N - (i * (i - 1)) / 2;
    int j = i + (kk - start);

    long row1 = (((long)b * N + i) * N + j) * C_DIM;
    long row2 = (((long)b * N + j) * N + i) * C_DIM;
    float a0=inp[row1+lane], a1=inp[row1+lane+32];
    float c0=inp[row2+lane], c1=inp[row2+lane+32];
    float mean1=warp_sum(a0+a1)*(1.0f/C_DIM);
    float mean2=warp_sum(c0+c1)*(1.0f/C_DIM);
    mean1=__shfl_sync(0xffffffff, mean1, 0);
    mean2=__shfl_sync(0xffffffff, mean2, 0);
    float d0=a0-mean1, d1=a1-mean1, e0=c0-mean2, e1=c1-mean2;
    float rstd1=rsqrtf(warp_sum(d0*d0+d1*d1)*(1.0f/C_DIM)+1.0e-5f);
    float rstd2=rsqrtf(warp_sum(e0*e0+e1*e1)*(1.0f/C_DIM)+1.0e-5f);
    rstd1=__shfl_sync(0xffffffff, rstd1, 0);
    rstd2=__shfl_sync(0xffffffff, rstd2, 0);
    float* s=sx+warp*C_DIM;
    s[lane] = 0.5f*(d0*rstd1*ln_w[lane]+ln_b[lane] + e0*rstd2*ln_w[lane]+ln_b[lane]);
    s[lane+32] = 0.5f*(d1*rstd1*ln_w[lane+32]+ln_b[lane+32] + e1*rstd2*ln_w[lane+32]+ln_b[lane+32]);
    __syncwarp();
    float acc=bias1[lane];
    #pragma unroll
    for(int k=0;k<C_DIM;k++) acc=fmaf(w1[lane*C_DIM+k], s[k], acc);
    float* y1=h1+warp*H_DIM;
    y1[lane]=acc>0.0f?acc:0.0f;
    __syncwarp();
    float acc2=bias2[lane];
    #pragma unroll
    for(int k=0;k<H_DIM;k++) acc2=fmaf(w2[lane*H_DIM+k], y1[k], acc2);
    float* y2=h2+warp*H_DIM;
    y2[lane]=acc2>0.0f?acc2:0.0f;
    __syncwarp();
    float zsum=warp_sum(w3[lane]*y2[lane]);
    if(lane==0){
        float p=1.0f/(1.0f+expf(-(zsum+bias3[0])));
        long o1=((long)b*N+i)*N+j;
        long o2=((long)b*N+j)*N+i;
        out[o1]=p;
        out[o2]=p;
    }
}

torch::Tensor 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=(int)pair_repr.size(0), N=(int)pair_repr.size(1);
    auto out=torch::empty({B,N,N}, pair_repr.options());
    int tri=N*(N+1)/2;
    int total_tri=B*tri;
    int blocks=(total_tri+WARPS-1)/WARPS;
    contact_kernel_tri<<<blocks,BLOCK,0,at::cuda::getCurrentCUDAStream()>>>(
        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>(), out.data_ptr<float>(), N, tri, total_tri);
    return out;
}
"""
_ext=load_inline(name="contact_head_ext_final", cpp_sources=cpp_source, cuda_sources=cuda_source,
                 functions=["contact_forward"], extra_cuda_cflags=["-O3","--use_fast_math"], extra_cflags=["-O3"], verbose=False)
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))
    def forward(self,pair_repr:torch.Tensor)->torch.Tensor:
        return _ext.contact_forward(pair_repr,self.layer_norm.weight,self.layer_norm.bias,
                                    self.mlp[0].weight,self.mlp[0].bias,self.mlp[2].weight,self.mlp[2].bias,self.mlp[4].weight,self.mlp[4].bias)