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


_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor pairwise_cuda(torch::Tensor coords);

torch::Tensor pairwise_forward(torch::Tensor coords) {
    return pairwise_cuda(coords);
}
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>

#define TILE 16

__global__ __launch_bounds__(256) void pairwise_kernel_3d(
    const float* __restrict__ coords,
    float* __restrict__ out,
    const int N)
{
    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int i = blockIdx.y * TILE + ty;
    const int j = blockIdx.x * TILE + tx;
    const int b = blockIdx.z;

    __shared__ float ci[TILE][3];
    __shared__ float cj[TILE][3];

    const int tid = ty * TILE + tx;
    if (tid < TILE * 3) {
        const int r = tid / 3;
        const int c = tid - r * 3;
        const int gi = blockIdx.y * TILE + r;
        const int gj = blockIdx.x * TILE + r;
        if (gi < N) ci[r][c] = __ldg(coords + ((b * N + gi) * 3 + c));
        if (gj < N) cj[r][c] = __ldg(coords + ((b * N + gj) * 3 + c));
    }
    __syncthreads();

    if (i < N && j < N) {
        const float dx = ci[ty][0] - cj[tx][0];
        const float dy = ci[ty][1] - cj[tx][1];
        const float dz = ci[ty][2] - cj[tx][2];
        const float s = fmaf(dx, dx, fmaf(dy, dy, dz * dz)) + 1.0e-8f;
        out[(static_cast<long long>(b) * N + i) * N + j] = sqrtf(s);
    }
}

torch::Tensor pairwise_cuda(torch::Tensor coords) {
    TORCH_CHECK(coords.is_cuda(), "coords must be a CUDA tensor");
    TORCH_CHECK(coords.scalar_type() == at::kFloat, "coords must be float32");
    TORCH_CHECK(coords.dim() == 3 && coords.size(2) == 3, "coords must have shape (B,N,3)");
    TORCH_CHECK(coords.is_contiguous(), "coords must be contiguous");

    const int B = static_cast<int>(coords.size(0));
    const int N = static_cast<int>(coords.size(1));
    auto out = torch::empty({B, N, N}, coords.options());

    const dim3 block(TILE, TILE, 1);
    const dim3 grid((N + TILE - 1) / TILE, (N + TILE - 1) / TILE, B);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    pairwise_kernel_3d<<<grid, block, 0, stream>>>(coords.data_ptr<float>(), out.data_ptr<float>(), N);
    return out;
}
"""

_ext = load_inline(
    name="pairwise_dist_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["pairwise_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, coords: torch.Tensor) -> torch.Tensor:
        return _ext.pairwise_forward(coords)