import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Compute all-pairs Euclidean distance matrix from a 3-D point cloud.
    A core primitive in molecular dynamics simulations, docking, and
    distance-geometry-based structure prediction.
    """

    def __init__(self):
        super().__init__()

    def forward(self, coords: torch.Tensor) -> torch.Tensor:
        """
        Args:
            coords: (B, N, 3) – 3-D coordinates of N atoms / residues
        Returns:
            dist: (B, N, N) – pairwise Euclidean distance matrix
        """
        diff = coords.unsqueeze(2) - coords.unsqueeze(1)  # (B, N, N, 3)
        dist = torch.sqrt((diff ** 2).sum(dim=-1) + 1e-8)
        return dist


num_atoms = 512
batch_size = 8


def get_inputs():
    p = popcorn_pri
    return [torch.randn(p.jitter_int(batch_size), p.jitter_int(num_atoms), 3)]


def get_init_inputs():
    return []

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>

#define TILE 16

__global__ void pairwise_dist_kernel(const float* __restrict__ coords, float* __restrict__ dist, int B, int N) {
    int b = blockIdx.z;
    int tile_i = blockIdx.y;
    int tile_j = blockIdx.x;
    
    int local_i = threadIdx.y;
    int local_j = threadIdx.x;
    
    int i = tile_i * TILE + local_i;
    int j = tile_j * TILE + local_j;
    
    float xi = 0, yi = 0, zi = 0;
    if (i < N) {
        xi = coords[b * N * 3 + i * 3 + 0];
        yi = coords[b * N * 3 + i * 3 + 1];
        zi = coords[b * N * 3 + i * 3 + 2];
    }
    
    float xj = 0, yj = 0, zj = 0;
    if (j < N) {
        xj = coords[b * N * 3 + j * 3 + 0];
        yj = coords[b * N * 3 + j * 3 + 1];
        zj = coords[b * N * 3 + j * 3 + 2];
    }
    
    float dx = xi - xj;
    float dy = yi - yj;
    float dz = zi - zj;
    float d2 = dx*dx + dy*dy + dz*dz;
    
    if (i < N && j < N) {
        dist[b * N * N + i * N + j] = sqrtf(d2 + 1e-8f);
    }
}

torch::Tensor pairwise_dist_cuda(torch::Tensor coords) {
    int B = coords.size(0);
    int N = coords.size(1);
    auto dist = torch::empty({B, N, N}, coords.options());
    
    dim3 blocks((N + TILE - 1) / TILE, (N + TILE - 1) / TILE, B);
    dim3 threads(TILE, TILE);
    
    pairwise_dist_kernel<<<blocks, threads>>>(
        coords.data_ptr<float>(),
        dist.data_ptr<float>(),
        B, N
    );
    return dist;
}
'''

cpp_source = '''
torch::Tensor pairwise_dist_cuda(torch::Tensor coords);
'''

module = load_inline(
    name='pairwise_dist',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['pairwise_dist_cuda'],
    verbose=True,
    extra_cuda_cflags=['-O3']
)

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.module = module
    
    def forward(self, coords):
        return self.module.pairwise_dist_cuda(coords)