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
from torch.utils.cpp_extension import load_inline
import math

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

__global__ void rbf_kernel(
    const float* __restrict__ distances,
    const float* __restrict__ offsets,
    float* __restrict__ output,
    float cutoff,
    float coeff,
    int64_t numel,
    int num_rbf
) {
    int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= numel) return;

    float d = distances[idx];
    float env = 0.0f;
    if (d < cutoff) {
        float arg = 3.141592653589793f * d / cutoff;
        env = 0.5f * (cosf(arg) + 1.0f);
    }
    for (int k = 0; k < num_rbf; ++k) {
        float diff = d - offsets[k];
        float rbf = expf(coeff * diff * diff);
        output[idx * num_rbf + k] = rbf * env;
    }
}

torch::Tensor rbf_cuda(torch::Tensor distances, torch::Tensor offsets, float cutoff, float coeff) {
    auto sizes = distances.sizes().vec();
    sizes.push_back(offsets.size(0));
    auto output = torch::empty(sizes, distances.options());

    int64_t numel = distances.numel();
    int num_rbf = offsets.size(0);

    const int block_size = 256;
    int num_blocks = (numel + block_size - 1) / block_size;

    rbf_kernel<<<num_blocks, block_size>>>(
        distances.data_ptr<float>(),
        offsets.data_ptr<float>(),
        output.data_ptr<float>(),
        cutoff,
        coeff,
        numel,
        num_rbf
    );

    return output;
}
"""

cpp_source = """
torch::Tensor rbf_cuda(torch::Tensor distances, torch::Tensor offsets, float cutoff, float coeff);
"""

module = load_inline(
    name="rbf_module",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["rbf_cuda"],
    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);
        coeff = -0.5 / (offsets[1] - offsets[0]).item() ** 2;
        self.coeff = coeff;

    def forward(self, distances: torch.Tensor) -> torch.Tensor:
        return module.rbf_cuda(distances, self.offsets, self.cutoff, self.coeff)