import torch
import torch.nn as nn
import math


class Model(nn.Module):
    """
    Radial Basis Function (RBF) expansion of interatomic distances.
    Standard featurization in molecular graph neural networks (SchNet,
    DimeNet, GemNet) that converts scalar distances into a feature vector
    of Gaussian-basis evaluations.
    """

    def __init__(self, num_rbf=64, cutoff=5.0):
        super().__init__()
        self.num_rbf = num_rbf
        self.cutoff = cutoff
        offsets = torch.linspace(0.0, cutoff, num_rbf)
        self.register_buffer("offsets", offsets)
        coeff = -0.5 / (offsets[1] - offsets[0]).item() ** 2
        self.coeff = coeff

    def forward(self, distances: torch.Tensor) -> torch.Tensor:
        """
        Args:
            distances: (B, N, N) – pairwise distance matrix (or (B, E) edge list)
        Returns:
            rbf_features: (*distances.shape, num_rbf) – expanded features
        """
        # Cosine cutoff envelope
        d_scaled = distances.unsqueeze(-1)  # (..., 1)
        envelope = 0.5 * (torch.cos(math.pi * d_scaled / self.cutoff) + 1.0)
        envelope = envelope * (d_scaled < self.cutoff).float()

        rbf = torch.exp(self.coeff * (d_scaled - self.offsets) ** 2)
        return rbf * envelope


num_atoms = 256
num_rbf = 64
cutoff = 5.0
batch_size = 8


def get_inputs():
    p = popcorn_pri
    coords = torch.randn(p.jitter_int(batch_size), p.jitter_int(num_atoms), 3)
    diff = coords.unsqueeze(2) - coords.unsqueeze(1)
    distances = torch.sqrt((diff ** 2).sum(-1) + 1e-8)
    return [distances]


def get_init_inputs():
    return [num_rbf, cutoff]

import torch
import torch.nn as nn
import math
from torch.utils.cpp_extension import load_inline

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

torch::Tensor rbf_forward(torch::Tensor distances, int64_t num_rbf, double cutoff);
"""

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

#ifndef CUDART_PI_F
#define CUDART_PI_F 3.141592654f
#endif

__global__ void rbf64_kernel(const float* __restrict__ dist,
                             float* __restrict__ out,
                             const int64_t D,
                             const float cutoff) {
    // Four warps per CTA; each warp expands one distance into 64 RBF values.
    const int warp = threadIdx.x >> 5;
    const int lane = threadIdx.x & 31;
    const int64_t idx = (int64_t)blockIdx.x * 4 + warp;
    if (idx >= D) return;

    float d = 0.0f;
    float env = 0.0f;
    if (lane == 0) {
        d = dist[idx];
        env = (d < cutoff) ? (0.5f * (cosf(CUDART_PI_F * d / cutoff) + 1.0f)) : 0.0f;
    }
    d = __shfl_sync(0xffffffff, d, 0);
    env = __shfl_sync(0xffffffff, env, 0);

    constexpr float inv63 = 1.0f / 63.0f;
    const float spacing = cutoff * inv63;
    const float coeff = -0.5f / (spacing * spacing);
    float* base = out + idx * 64;

    int j = lane;
    float x = d - (float)j * spacing;
    base[j] = __expf(coeff * x * x) * env;

    j = lane + 32;
    x = d - (float)j * spacing;
    base[j] = __expf(coeff * x * x) * env;
}

__global__ void rbf_generic_kernel(const float* __restrict__ dist,
                                   float* __restrict__ out,
                                   const int64_t total,
                                   const int64_t num_rbf,
                                   const float cutoff) {
    const int64_t k = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (k >= total) return;
    const int64_t d_idx = k / num_rbf;
    const int64_t j = k - d_idx * num_rbf;
    const float d = dist[d_idx];
    const float spacing = cutoff / (float)(num_rbf - 1);
    const float coeff = -0.5f / (spacing * spacing);
    const float env = (d < cutoff) ? (0.5f * (cosf(CUDART_PI_F * d / cutoff) + 1.0f)) : 0.0f;
    const float off = (float)j * spacing;
    const float x = d - off;
    out[k] = __expf(coeff * x * x) * env;
}

torch::Tensor rbf_forward(torch::Tensor distances, int64_t num_rbf, double cutoff_d) {
    TORCH_CHECK(distances.is_cuda(), "distances must be a CUDA tensor");
    TORCH_CHECK(distances.scalar_type() == at::kFloat, "only float32 distances are supported");
    TORCH_CHECK(num_rbf > 1, "num_rbf must be greater than 1");

    std::vector<int64_t> shape = distances.sizes().vec();
    shape.push_back(num_rbf);
    auto out = torch::empty(shape, distances.options());

    const int64_t D = distances.numel();
    if (D == 0) return out;

    const float cutoff = static_cast<float>(cutoff_d);
    const float* dist_ptr = distances.data_ptr<float>();
    float* out_ptr = out.data_ptr<float>();
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (num_rbf == 64) {
        const int threads = 128;
        const dim3 blocks((unsigned int)((D + 3) / 4));
        rbf64_kernel<<<blocks, threads, 0, stream>>>(dist_ptr, out_ptr, D, cutoff);
    } else {
        const int threads = 256;
        const int64_t total = D * num_rbf;
        const dim3 blocks((unsigned int)((total + threads - 1) / threads));
        rbf_generic_kernel<<<blocks, threads, 0, stream>>>(dist_ptr, out_ptr, total, num_rbf, cutoff);
    }
    return out;
}
"""

_rbf_ext = load_inline(
    name="rbf_expansion_inline_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["rbf_forward"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, num_rbf=64, cutoff=5.0):
        super().__init__()
        self.num_rbf = num_rbf
        self.cutoff = cutoff
        offsets = torch.linspace(0.0, cutoff, num_rbf)
        self.register_buffer("offsets", offsets)
        self.coeff = -0.5 / (offsets[1] - offsets[0]).item() ** 2

    def forward(self, distances: torch.Tensor) -> torch.Tensor:
        if not distances.is_contiguous():
            distances = distances.contiguous()
        return _rbf_ext.rbf_forward(distances, int(self.num_rbf), float(self.cutoff))