import torch
import torch.nn as nn


class Model(nn.Module):
    """
    RBF (Squared Exponential) covariance matrix computation for Gaussian
    Processes.  Given two sets of input points, computes the full
    kernel matrix K[i,j] = sigma^2 * exp(-||x_i - x_j||^2 / (2 * l^2)).
    The dominant cost in GP inference and a natural target for fusion.
    """

    def __init__(self, input_dim, lengthscale=1.0, variance=1.0):
        super().__init__()
        self.log_lengthscale = nn.Parameter(torch.tensor(lengthscale).log())
        self.log_variance = nn.Parameter(torch.tensor(variance).log())

    def forward(self, X1: torch.Tensor, X2: torch.Tensor) -> torch.Tensor:
        """
        Args:
            X1: (B, N, D) – first set of inputs
            X2: (B, M, D) – second set of inputs
        Returns:
            K: (B, N, M) – covariance matrix
        """
        l2 = self.log_lengthscale.exp() ** 2
        var = self.log_variance.exp()

        # Squared distances via expansion
        X1_sq = (X1 ** 2).sum(-1, keepdim=True)  # (B, N, 1)
        X2_sq = (X2 ** 2).sum(-1, keepdim=True)  # (B, M, 1)
        dist_sq = X1_sq - 2.0 * X1 @ X2.transpose(-1, -2) + X2_sq.transpose(-1, -2)
        dist_sq = dist_sq.clamp(min=0.0)

        return var * torch.exp(-0.5 * dist_sq / l2)


input_dim = 8
num_train = 512
num_test = 128
batch_size = 16


def get_inputs():
    X1 = torch.randn(popcorn_pri.jitter_int(batch_size), popcorn_pri.jitter_int(num_train), input_dim)
    X2 = torch.randn(popcorn_pri.jitter_int(batch_size), popcorn_pri.jitter_int(num_test), input_dim)
    return [X1, X2]


def get_init_inputs():
    return [input_dim]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

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

__global__ void rbf_kernel(
    const float* __restrict__ X1,
    const float* __restrict__ X2,
    float* __restrict__ K,
    const float var,
    const float inv_2l2,
    const int B, const int N, const int M, const int D
) {
    int b = blockIdx.z;
    int i = blockIdx.y * blockDim.y + threadIdx.y;
    int j = blockIdx.x * blockDim.x + threadIdx.x;

    if (i >= N || j >= M || b >= B) return;

    float dist_sq = 0.0f;
    const float* x1_ptr = X1 + (b * N + i) * D;
    const float* x2_ptr = X2 + (b * M + j) * D;

    for (int d = 0; d < D; ++d) {
        float diff = x1_ptr[d] - x2_ptr[d];
        dist_sq += diff * diff;
    }

    float val = var * expf(-0.5f * dist_sq * inv_2l2);
    K[(b * N + i) * M + j] = val;
}

torch::Tensor rbf_cuda(torch::Tensor X1, torch::Tensor X2, float var, float inv_2l2) {
    const int B = X1.size(0);
    const int N = X1.size(1);
    const int D = X1.size(2);
    const int M = X2.size(1);

    auto K = torch::empty({B, N, M}, X1.options());

    dim3 threads(16, 16);
    dim3 blocks((M + 15) / 16, (N + 15) / 16, B);

    rbf_kernel<<<blocks, threads>>>(
        X1.data_ptr<float>(),
        X2.data_ptr<float>(),
        K.data_ptr<float>(),
        var,
        inv_2l2,
        B, N, M, D
    );

    return K;
}
'''

cpp_source = r'''
torch::Tensor rbf_cuda(torch::Tensor X1, torch::Tensor X2, float var, float inv_2l2);
'''

rbf_module = load_inline(
    name="rbf_cuda",
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=["rbf_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"]
)

class ModelNew(nn.Module):
    def __init__(self, input_dim, lengthscale=1.0, variance=1.0):
        super().__init__()
        self.log_lengthscale = nn.Parameter(torch.tensor(lengthscale).log())
        self.log_variance = nn.Parameter(torch.tensor(variance).log())

    def forward(self, X1: torch.Tensor, X2: torch.Tensor) -> torch.Tensor:
        l2 = self.log_lengthscale.exp() ** 2
        var = self.log_variance.exp()
        inv_2l2 = 1.0 / (2.0 * l2)
        return rbf_module.rbf_cuda(X1, X2, var.item(), inv_2l2.item())