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 = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

// RBF kernel with precomputed squared norms in shared memory
// Each thread computes one output element

template <typename scalar_t, int TILE_N, int TILE_M, int D>
__global__ void __launch_bounds__(1024) rbf_kernel_forward(
    const scalar_t* __restrict__ X1,  // (B, N, D)
    const scalar_t* __restrict__ X2,  // (B, M, D)
    scalar_t* __restrict__ K,          // (B, N, M)
    const scalar_t lengthscale_sq,    // l^2
    const scalar_t variance,          // sigma^2
    const int B, const int N, const int M
) {
    const int batch = blockIdx.z;
    const int tile_i_start = blockIdx.y * TILE_N;
    const int tile_j_start = blockIdx.x * TILE_M;
    
    // Shared memory: X1/X2 vectors and their squared norms
    // Pad to avoid bank conflicts
    __shared__ scalar_t X1_tile[TILE_N][D + 4];
    __shared__ scalar_t X2_tile[TILE_M][D + 4];
    __shared__ scalar_t X1_sq_cache[TILE_N];
    __shared__ scalar_t X2_sq_cache[TILE_M];
    
    const int tid = threadIdx.x;
    const int block_size = blockDim.x;
    
    // Cooperative load of X1 tile and compute squared norms
    // Each thread loads one element, multiple rounds if needed
    const int x1_elements = TILE_N * D;
    for (int idx = tid; idx < x1_elements; idx += block_size) {
        const int row = idx / D;
        const int col = idx % D;
        const int global_row = tile_i_start + row;
        scalar_t val = static_cast<scalar_t>(0);
        if (global_row < N) {
            val = X1[batch * N * D + global_row * D + col];
        }
        X1_tile[row][col] = val;
    }
    
    // Cooperative load of X2 tile
    const int x2_elements = TILE_M * D;
    for (int idx = tid; idx < x2_elements; idx += block_size) {
        const int row = idx / D;
        const int col = idx % D;
        const int global_row = tile_j_start + row;
        scalar_t val = static_cast<scalar_t>(0);
        if (global_row < M) {
            val = X2[batch * M * D + global_row * D + col];
        }
        X2_tile[row][col] = val;
    }
    
    __syncthreads();
    
    // Compute squared norms for X1 rows
    // Distribute across threads
    for (int row = tid; row < TILE_N; row += block_size) {
        scalar_t sq = static_cast<scalar_t>(0);
        for (int d = 0; d < D; d++) {
            sq += X1_tile[row][d] * X1_tile[row][d];
        }
        X1_sq_cache[row] = sq;
    }
    
    // Compute squared norms for X2 rows
    for (int row = tid; row < TILE_M; row += block_size) {
        scalar_t sq = static_cast<scalar_t>(0);
        for (int d = 0; d < D; d++) {
            sq += X2_tile[row][d] * X2_tile[row][d];
        }
        X2_sq_cache[row] = sq;
    }
    
    __syncthreads();
    
    // Each thread computes one output element
    const int total_outputs = TILE_N * TILE_M;
    
    for (int out_idx = tid; out_idx < total_outputs; out_idx += block_size) {
        const int lane_i = out_idx / TILE_M;
        const int lane_j = out_idx % TILE_M;
        
        const int i = tile_i_start + lane_i;
        const int j = tile_j_start + lane_j;
        
        if (i < N && j < M) {
            // Get precomputed squared norms
            scalar_t x1_sq = X1_sq_cache[lane_i];
            scalar_t x2_sq = X2_sq_cache[lane_j];
            
            // Compute dot product
            scalar_t dot = static_cast<scalar_t>(0);
            for (int d = 0; d < D; d++) {
                dot += X1_tile[lane_i][d] * X2_tile[lane_j][d];
            }
            
            // Squared distance
            scalar_t dist_sq = x1_sq - static_cast<scalar_t>(2) * dot + x2_sq;
            dist_sq = dist_sq > static_cast<scalar_t>(0) ? dist_sq : static_cast<scalar_t>(0);
            
            // RBF kernel
            scalar_t exp_arg = static_cast<scalar_t>(-0.5) * dist_sq / lengthscale_sq;
            K[batch * N * M + i * M + j] = variance * exp(exp_arg);
        }
    }
}

torch::Tensor rbf_forward_cuda(
    torch::Tensor X1,
    torch::Tensor X2,
    double lengthscale_sq,
    double variance
) {
    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());
    
    const int TILE_N = 32;
    const int TILE_M = 32;
    const int BLOCK_SIZE = 1024;
    
    const dim3 grid((M + TILE_M - 1) / TILE_M, (N + TILE_N - 1) / TILE_N, B);
    const dim3 block(BLOCK_SIZE);
    
    AT_DISPATCH_FLOATING_TYPES(X1.scalar_type(), "rbf_kernel_forward", ([&] {
        rbf_kernel_forward<scalar_t, TILE_N, TILE_M, 8><<<grid, block>>>(
            X1.data_ptr<scalar_t>(),
            X2.data_ptr<scalar_t>(),
            K.data_ptr<scalar_t>(),
            static_cast<scalar_t>(lengthscale_sq),
            static_cast<scalar_t>(variance),
            B, N, M
        );
    }));
    
    return K;
}
"""

cpp_source = """
torch::Tensor rbf_forward_cuda(torch::Tensor X1, torch::Tensor X2, double lengthscale_sq, double variance);
"""

rbf_module = load_inline(
    name='rbf_cuda_v3',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['rbf_forward_cuda'],
    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())
        self.input_dim = input_dim

    def forward(self, X1: torch.Tensor, X2: torch.Tensor) -> torch.Tensor:
        l2 = self.log_lengthscale.exp() ** 2
        var = self.log_variance.exp()
        
        # Ensure inputs are contiguous
        X1 = X1.contiguous()
        X2 = X2.contiguous()
        
        return rbf_module.rbf_forward_cuda(X1, X2, l2.item(), var.item())