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>
#include <math.h>

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

    int elem_idx = idx / num_rbf;
    int rbf_idx = idx % num_rbf;

    float d = distances[elem_idx];
    float off = offsets[rbf_idx];
    float d_scaled = d;

    float envelope = 0.5f * (cosf(M_PI * d_scaled / cutoff) + 1.0f);
    envelope *= (d_scaled < cutoff) ? 1.0f : 0.0f;

    float rbf_val = expf(coeff * (d_scaled - off) * (d_scaled - off));
    output[idx] = rbf_val * envelope;
}

torch::Tensor rbf_cuda(torch::Tensor distances, torch::Tensor offsets, float cutoff, float coeff) {
    int total_elements = distances.numel();
    int num_rbf = offsets.size(0);
    auto output = torch::empty({total_elements, num_rbf}, distances.options());
    
    int total_out = total_elements * num_rbf;
    int threads = 256;
    int blocks = (total_out + threads - 1) / threads;

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

    // reshape to original shape + num_rbf
    std::vector<int64_t> new_shape(distances.sizes().begin(), distances.sizes().end());
    new_shape.push_back(num_rbf);
    return output.view(new_shape);
}
'''

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

rbf_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 rbf_module.rbf_cuda(distances, self.offsets, self.cutoff, self.coeff)