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

cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math_constants.h>

__global__ __launch_bounds__(256, 4)
void rbf64_pair_kernel(const float* __restrict__ distances,
                       float* __restrict__ out,
                       long long n, float cutoff, float coeff) {
    int tid = threadIdx.x;
    int lane = tid & 31;
    int warp = tid >> 5;
    long long i = (long long)blockIdx.x * 8 + warp;

    float d = 0.0f;
    float env = 0.0f;
    if (i < n && lane == 0) {
        d = distances[i];
        if (d < cutoff) {
            env = 0.5f * (cosf(CUDART_PI_F * d / cutoff) + 1.0f);
        }
    }
    d = __shfl_sync(0xffffffff, d, 0);
    env = __shfl_sync(0xffffffff, env, 0);

    if (i < n) {
        float cstep = cutoff * (1.0f / 63.0f);
        int k = lane * 2;
        float x0 = d - ((float)k) * cstep;
        float x1 = d - ((float)(k + 1)) * cstep;
        float2 v;
        v.x = expf(coeff * x0 * x0) * env;
        v.y = expf(coeff * x1 * x1) * env;
        reinterpret_cast<float2*>(out + i * 64)[lane] = v;
    }
}

__global__ __launch_bounds__(256, 4)
void rbf_generic_kernel(const float* __restrict__ distances,
                        float* __restrict__ out,
                        long long total, int m, float cutoff, float coeff) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < total) {
        int k = (int)(idx % m);
        long long i = idx / m;
        float d = distances[i];
        float env = 0.0f;
        if (d < cutoff) {
            env = 0.5f * (cosf(CUDART_PI_F * d / cutoff) + 1.0f);
        }
        float center = ((float)k) * cutoff / ((float)(m - 1));
        float x = d - center;
        out[idx] = expf(coeff * x * x) * env;
    }
}

torch::Tensor rbf_forward(torch::Tensor distances,
                          int64_t num_rbf,
                          double cutoff_d,
                          double coeff_d) {
    auto out = torch::empty({distances.numel(), num_rbf}, distances.options());
    long long n = distances.numel();
    float cutoff = (float)cutoff_d;
    float coeff = (float)coeff_d;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (num_rbf == 64) {
        rbf64_pair_kernel<<<(unsigned int)((n + 7) / 8), 256, 0, stream>>>(
            distances.data_ptr<float>(), out.data_ptr<float>(), n, cutoff, coeff);
    } else {
        long long total = n * num_rbf;
        rbf_generic_kernel<<<(unsigned int)((total + 255) / 256), 256, 0, stream>>>(
            distances.data_ptr<float>(), out.data_ptr<float>(), total, (int)num_rbf, cutoff, coeff);
    }
    return out;
}
'''

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

_rbf_ext = load_inline(
    name="rbf_ext_final_pair",
    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)
        coeff = -0.5 / (offsets[1] - offsets[0]).item() ** 2
        self.coeff = coeff

    def forward(self, distances: torch.Tensor) -> torch.Tensor:
        out = _rbf_ext.rbf_forward(distances, int(self.num_rbf), float(self.cutoff), float(self.coeff))
        return out.view(*distances.shape, self.num_rbf)