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

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

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

    float d = distances[idx];
    float d_scaled = d;
    bool mask = d_scaled < cutoff;
    float env = 0.0f;
    if (mask) {
        float arg = 3.141592653589793f * d_scaled / cutoff;
        env = 0.5f * (cosf(arg) + 1.0f);
    }

    float* out_ptr = output + idx * num_rbf;
    for (int i = 0; i < num_rbf; i++) {
        float diff = d_scaled - offsets[i];
        float rbf_val = expf(coeff * diff * diff);
        out_ptr[i] = rbf_val * env;
    }
}

torch::Tensor rbf_expansion_cuda(torch::Tensor distances, torch::Tensor offsets, float coeff, float cutoff, int num_rbf) {
    auto sizes = distances.sizes();
    std::vector<int64_t> out_sizes(sizes.begin(), sizes.end());
    out_sizes.push_back(num_rbf);
    auto output = torch::empty(out_sizes, distances.options());

    int total_elements = distances.numel();
    int threads = 256;
    int blocks = (total_elements + threads - 1) / threads;

    rbf_kernel<<<blocks, threads>>>(
        distances.data_ptr<float>(),
        output.data_ptr<float>(),
        offsets.data_ptr<float>(),
        coeff,
        cutoff,
        num_rbf,
        total_elements
    );

    return output;
}
'''

cpp_source = '''
torch::Tensor rbf_expansion_cuda(torch::Tensor distances, torch::Tensor offsets, float coeff, float cutoff, int num_rbf);
'''

rbf_module = load_inline(
    name='rbf_cuda',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['rbf_expansion_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 rbf_module.rbf_expansion_cuda(distances, self.offsets, self.coeff, self.cutoff, self.num_rbf)