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 <cmath>

constexpr float PI = 3.14159265358979323846f;

// Vectorized kernel with float4 for better memory throughput
__global__ void rbf_kernel_v3(
    const float* __restrict__ distances,
    const float* __restrict__ offsets,
    float* __restrict__ output,
    const int num_elements,
    const int num_rbf,
    const float cutoff,
    const float coeff
) {
    // Shared memory for offsets
    __shared__ float s_offsets[64];
    
    // Cooperatively load offsets
    for (int i = threadIdx.x; i < num_rbf; i += blockDim.x) {
        s_offsets[i] = offsets[i];
    }
    __syncthreads();
    
    const int vec_size = 4;
    int idx = (blockIdx.x * blockDim.x + threadIdx.x) * vec_size;
    
    if (idx >= num_elements) return;
    
    // Load 4 distances at once (if aligned and within bounds)
    float d[4];
    bool valid[4];
    
    #pragma unroll
    for (int i = 0; i < vec_size; i++) {
        int cur_idx = idx + i;
        valid[i] = cur_idx < num_elements;
        d[i] = valid[i] ? distances[cur_idx] : 0.0f;
    }
    
    // Process each distance
    for (int i = 0; i < vec_size; i++) {
        if (!valid[i]) continue;
        
        float envelope = 0.5f * (cosf(PI * d[i] / cutoff) + 1.0f);
        if (d[i] >= cutoff) envelope = 0.0f;
        
        // Process all RBF values
        for (int r = 0; r < num_rbf; r++) {
            float offset = s_offsets[r];
            float diff = d[i] - offset;
            float rbf = expf(coeff * diff * diff);
            output[(idx + i) * num_rbf + r] = rbf * envelope;
        }
    }
}

torch::Tensor rbf_forward_cuda_v3(
    torch::Tensor distances,
    torch::Tensor offsets,
    float cutoff,
    float coeff
) {
    int num_elements = distances.numel();
    int num_rbf = offsets.size(0);
    
    std::vector<int64_t> output_shape = distances.sizes().vec();
    output_shape.push_back(num_rbf);
    auto output = torch::empty(output_shape, distances.options());
    
    const int vec_size = 4;
    int threads = 256;
    int blocks = (num_elements + threads * vec_size - 1) / (threads * vec_size);
    
    rbf_kernel_v3<<<blocks, threads>>>(
        distances.data_ptr<float>(),
        offsets.data_ptr<float>(),
        output.data_ptr<float>(),
        num_elements,
        num_rbf,
        cutoff,
        coeff
    );
    
    return output;
}
"""

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

rbf_extension_v3 = load_inline(
    name='rbf_extension_v3',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['rbf_forward_cuda_v3'],
    extra_cuda_cflags=['-O3', '-use_fast_math'],
    verbose=True
)


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:
        distances = distances.contiguous()
        offsets = self.offsets.contiguous()
        return rbf_extension_v3.rbf_forward_cuda_v3(distances, offsets, self.cutoff, self.coeff)