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

_cpp_src = r'''
#include <torch/extension.h>
torch::Tensor rbf_forward(torch::Tensor X1, torch::Tensor X2, torch::Tensor log_l, torch::Tensor log_v);
'''

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

#define TILE_N 8
#define TILE_M 32
#define D_FIXED 8
#define DPAD 9

__global__ void rbf_d8_tiled_kernel(const float* __restrict__ X1,
                                   const float* __restrict__ X2,
                                   const float* __restrict__ log_l,
                                   const float* __restrict__ log_v,
                                   float* __restrict__ O,
                                   int B, int N, int M) {
    __shared__ float sx1[TILE_N][DPAD];
    __shared__ float sx2[TILE_M][DPAD];
    __shared__ float scale;
    __shared__ float var;

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int tid = ty * TILE_M + tx;

    if (tid == 0) {
        float l = __expf(log_l[0]);
        scale = -0.5f / (l * l);
        var = __expf(log_v[0]);
    }

    const int b = blockIdx.z;
    const int n0 = blockIdx.y * TILE_N;
    const int m0 = blockIdx.x * TILE_M;

    if (tid < TILE_N * D_FIXED) {
        int row = tid >> 3;
        int d = tid & 7;
        int n = n0 + row;
        sx1[row][d] = (n < N) ? X1[((b * N + n) * D_FIXED + d)] : 0.0f;
    }
    if (tid < TILE_M * D_FIXED) {
        int row = tid >> 3;
        int d = tid & 7;
        int m = m0 + row;
        sx2[row][d] = (m < M) ? X2[((b * M + m) * D_FIXED + d)] : 0.0f;
    }
    __syncthreads();

    const int n = n0 + ty;
    const int m = m0 + tx;
    if (n < N && m < M) {
        float s = 0.0f;
        #pragma unroll
        for (int d = 0; d < D_FIXED; ++d) {
            float diff = sx1[ty][d] - sx2[tx][d];
            s = fmaf(diff, diff, s);
        }
        O[(b * N + n) * M + m] = var * __expf(scale * s);
    }
}

__global__ void rbf_generic_kernel(const float* __restrict__ X1,
                                   const float* __restrict__ X2,
                                   const float* __restrict__ log_l,
                                   const float* __restrict__ log_v,
                                   float* __restrict__ O,
                                   int total, int N, int M, int D) {
    __shared__ float scale;
    __shared__ float var;
    if (threadIdx.x == 0) {
        float l = __expf(log_l[0]);
        scale = -0.5f / (l * l);
        var = __expf(log_v[0]);
    }
    __syncthreads();
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;
    int m = idx % M;
    int t = idx / M;
    int n = t % N;
    int b = t / N;
    const float* p1 = X1 + ((b * N + n) * D);
    const float* p2 = X2 + ((b * M + m) * D);
    float s = 0.0f;
    for (int d = 0; d < D; ++d) {
        float diff = p1[d] - p2[d];
        s = fmaf(diff, diff, s);
    }
    O[idx] = var * __expf(scale * s);
}

torch::Tensor rbf_forward(torch::Tensor X1, torch::Tensor X2, torch::Tensor log_l, torch::Tensor log_v) {
    const int B = (int)X1.size(0);
    const int N = (int)X1.size(1);
    const int D = (int)X1.size(2);
    const int M = (int)X2.size(1);
    auto O = torch::empty({B, N, M}, X1.options());
    if (D == 8) {
        dim3 block(TILE_M, TILE_N, 1);
        dim3 grid((M + TILE_M - 1) / TILE_M, (N + TILE_N - 1) / TILE_N, B);
        rbf_d8_tiled_kernel<<<grid, block>>>(X1.data_ptr<float>(), X2.data_ptr<float>(),
                                             log_l.data_ptr<float>(), log_v.data_ptr<float>(),
                                             O.data_ptr<float>(), B, N, M);
    } else {
        int total = B * N * M;
        int threads = 256;
        int blocks = (total + threads - 1) / threads;
        rbf_generic_kernel<<<blocks, threads>>>(X1.data_ptr<float>(), X2.data_ptr<float>(),
                                                log_l.data_ptr<float>(), log_v.data_ptr<float>(),
                                                O.data_ptr<float>(), total, N, M, D);
    }
    return O;
}
'''

_ext = load_inline(name='rbf_cov_ext_final', cpp_sources=_cpp_src, cuda_sources=_cuda_src,
                   functions=['rbf_forward'], verbose=False,
                   extra_cuda_cflags=['-O3', '--use_fast_math'])

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:
        return _ext.rbf_forward(X1, X2, self.log_lengthscale, self.log_variance)