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

# Fully fused kernel - compute everything in CUDA
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

// Fully fused RBF kernel - computes squared norms internally
template<int TILE_SIZE, int D_MAX>
__global__ void rbf_kernel_fused(
    const float* __restrict__ X1,      // (B, N, D)
    const float* __restrict__ X2,      // (B, M, D)
    float* __restrict__ output,         // (B, N, M)
    float neg_half_over_l2,             // -0.5 / l2
    float var,                          // variance
    int B, int N, int M, int D
) {
    // Shared memory for tile of X1 rows and X2 rows
    __shared__ float s_X1[TILE_SIZE * D_MAX];
    __shared__ float s_X2[TILE_SIZE * D_MAX];
    __shared__ float s_X1_sq[TILE_SIZE];
    __shared__ float s_X2_sq[TILE_SIZE];
    
    const int batch_idx = blockIdx.z;
    const int tile_n_start = blockIdx.y * TILE_SIZE;  // starting row in X1
    const int tile_m_start = blockIdx.x * TILE_SIZE;  // starting col in X2
    
    const int tid = threadIdx.x;
    const int block_size = blockDim.x;
    
    // Load TILE_SIZE rows of X1 into shared memory
    for (int i = tid; i < TILE_SIZE * D; i += block_size) {
        int row = i / D;
        int d = i % D;
        int global_row = tile_n_start + row;
        if (global_row < N) {
            float val = X1[batch_idx * N * D + global_row * D + d];
            s_X1[i] = val;
        } else {
            s_X1[i] = 0.0f;
        }
    }
    
    // Load TILE_SIZE rows of X2 into shared memory
    for (int i = tid; i < TILE_SIZE * D; i += block_size) {
        int row = i / D;
        int d = i % D;
        int global_row = tile_m_start + row;
        if (global_row < M) {
            float val = X2[batch_idx * M * D + global_row * D + d];
            s_X2[i] = val;
        } else {
            s_X2[i] = 0.0f;
        }
    }
    
    __syncthreads();
    
    // Compute squared norms for X1 tile rows
    for (int row = tid; row < TILE_SIZE; row += block_size) {
        float sum = 0.0f;
        for (int d = 0; d < D; d++) {
            float val = s_X1[row * D + d];
            sum += val * val;
        }
        s_X1_sq[row] = sum;
    }
    
    // Compute squared norms for X2 tile rows
    for (int row = tid; row < TILE_SIZE; row += block_size) {
        float sum = 0.0f;
        for (int d = 0; d < D; d++) {
            float val = s_X2[row * D + d];
            sum += val * val;
        }
        s_X2_sq[row] = sum;
    }
    
    __syncthreads();
    
    // Each thread computes one output element
    const int local_n = tid / TILE_SIZE;  // row within tile
    const int local_m = tid % TILE_SIZE;  // col within tile
    
    const int global_n = tile_n_start + local_n;
    const int global_m = tile_m_start + local_m;
    
    if (global_n < N && global_m < M) {
        // Compute dot product
        float dot = 0.0f;
        for (int d = 0; d < D; d++) {
            dot += s_X1[local_n * D + d] * s_X2[local_m * D + d];
        }
        
        // Get squared norms
        float x1_sq = s_X1_sq[local_n];
        float x2_sq = s_X2_sq[local_m];
        
        // Compute squared distance
        float dist_sq = x1_sq - 2.0f * dot + x2_sq;
        
        // Clamp to 0 (handle numerical errors)
        dist_sq = fmaxf(dist_sq, 0.0f);
        
        // Compute RBF
        float result = var * expf(dist_sq * neg_half_over_l2);
        output[batch_idx * N * M + global_n * M + global_m] = result;
    }
}

torch::Tensor rbf_forward(
    torch::Tensor X1,
    torch::Tensor X2,
    float neg_half_over_l2,
    float var
) {
    TORCH_CHECK(X1.is_cuda() && X2.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(X1.is_contiguous() && X2.is_contiguous(), "Inputs must be contiguous");
    
    int B = X1.size(0);
    int N = X1.size(1);
    int D = X1.size(2);
    int M = X2.size(1);
    
    auto output = torch::empty({B, N, M}, X1.options());
    
    const int TILE_SIZE = 16;
    const int BLOCK_SIZE = TILE_SIZE * TILE_SIZE;  // 256 threads
    
    dim3 grid((M + TILE_SIZE - 1) / TILE_SIZE, (N + TILE_SIZE - 1) / TILE_SIZE, B);
    dim3 block(BLOCK_SIZE);
    
    rbf_kernel_fused<TILE_SIZE, 16><<<grid, block>>>(
        X1.data_ptr<float>(),
        X2.data_ptr<float>(),
        output.data_ptr<float>(),
        neg_half_over_l2,
        var,
        B, N, M, D
    );
    
    return output;
}
"""

cpp_source = """
torch::Tensor rbf_forward(
    torch::Tensor X1,
    torch::Tensor X2,
    float neg_half_over_l2,
    float var
);
"""

rbf_extension = load_inline(
    name='rbf_extension',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['rbf_forward'],
    verbose=True
)


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()
        
        # Call fully fused CUDA kernel
        neg_half_over_l2 = -0.5 / l2.item()
        
        return rbf_extension.rbf_forward(
            X1.contiguous(),
            X2.contiguous(),
            neg_half_over_l2,
            var.item()
        )