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

// Shared memory tiled kernel for pairwise distances
template <int TILE_SIZE>
__global__ void pairwise_distance_kernel(
    const float* __restrict__ coords,
    float* __restrict__ dist,
    const int B,
    const int N
) {
    int b = blockIdx.z;
    int i_start = blockIdx.y * TILE_SIZE;
    int j_start = blockIdx.x * TILE_SIZE;
    
    int i_local = threadIdx.y;
    int j_local = threadIdx.x;
    
    int i = i_start + i_local;
    int j = j_start + j_local;
    
    // Shared memory for tiles of coordinates
    __shared__ float s_i_coords[TILE_SIZE][3];
    __shared__ float s_j_coords[TILE_SIZE][3];
    
    // Load coordinates for i points
    if (i < N) {
        s_i_coords[i_local][0] = coords[b * N * 3 + i * 3 + 0];
        s_i_coords[i_local][1] = coords[b * N * 3 + i * 3 + 1];
        s_i_coords[i_local][2] = coords[b * N * 3 + i * 3 + 2];
    }
    
    // Load coordinates for j points
    if (j < N) {
        s_j_coords[j_local][0] = coords[b * N * 3 + j * 3 + 0];
        s_j_coords[j_local][1] = coords[b * N * 3 + j * 3 + 1];
        s_j_coords[j_local][2] = coords[b * N * 3 + j * 3 + 2];
    }
    
    __syncthreads();
    
    // Compute distance
    if (i < N && j < N) {
        float dx = s_i_coords[i_local][0] - s_j_coords[j_local][0];
        float dy = s_i_coords[i_local][1] - s_j_coords[j_local][1];
        float dz = s_i_coords[i_local][2] - s_j_coords[j_local][2];
        float dist_sq = dx * dx + dy * dy + dz * dz;
        dist[b * N * N + i * N + j] = sqrtf(dist_sq + 1e-8f);
    }
}

void launch_pairwise_distance(const float* coords, float* dist, int B, int N) {
    const int TILE_SIZE = 16;
    dim3 threadsPerBlock(TILE_SIZE, TILE_SIZE);
    dim3 numBlocks((N + TILE_SIZE - 1) / TILE_SIZE,
                   (N + TILE_SIZE - 1) / TILE_SIZE,
                   B);
    pairwise_distance_kernel<TILE_SIZE><<<numBlocks, threadsPerBlock>>>(coords, dist, B, N);
}
"""

cpp_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

void launch_pairwise_distance(const float* coords, float* dist, int B, int N);

torch::Tensor pairwise_distance_cuda(torch::Tensor coords) {
    TORCH_CHECK(coords.dim() == 3, "coords must be 3D");
    TORCH_CHECK(coords.size(2) == 3, "coords last dimension must be 3");
    
    int B = coords.size(0);
    int N = coords.size(1);
    auto dist = torch::empty({B, N, N}, coords.options());
    
    const float* coords_ptr = coords.data_ptr<float>();
    float* dist_ptr = dist.data_ptr<float>();
    
    launch_pairwise_distance(coords_ptr, dist_ptr, B, N);
    
    return dist;
}
"""

module = load_inline(
    name='pairwise_distance_ext',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['pairwise_distance_cuda'],
    extra_cuda_cflags=['-O3'],
    verbose=False
)

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

    def forward(self, coords: torch.Tensor) -> torch.Tensor:
        return module.pairwise_distance_cuda(coords.contiguous())