import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Gaussian Process posterior predictive distribution via Cholesky
    decomposition.  Computes the predictive mean and variance at test
    points given training data and an RBF kernel.  The Cholesky solve
    is the computational bottleneck in exact GP inference.
    """

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

    def _rbf_kernel(self, X1, X2):
        l2 = self.log_lengthscale.exp() ** 2
        var = self.log_variance.exp()
        dist_sq = ((X1.unsqueeze(2) - X2.unsqueeze(1)) ** 2).sum(-1)
        return var * torch.exp(-0.5 * dist_sq / l2)

    def forward(
        self, X_train: torch.Tensor, Y_train: torch.Tensor, X_test: torch.Tensor
    ) -> tuple:
        """
        Args:
            X_train: (B, N, D)
            Y_train: (B, N, 1)
            X_test:  (B, M, D)
        Returns:
            (mean, var):
                mean: (B, M, 1) – predictive mean
                var:  (B, M)    – predictive variance
        """
        K_nn = self._rbf_kernel(X_train, X_train)
        K_nn = K_nn + self.noise * torch.eye(K_nn.shape[-1], device=K_nn.device).unsqueeze(0)
        K_nm = self._rbf_kernel(X_train, X_test)
        K_mm = self._rbf_kernel(X_test, X_test)

        L = torch.linalg.cholesky(K_nn)
        alpha = torch.cholesky_solve(Y_train, L)

        mean = K_nm.transpose(-1, -2) @ alpha

        v = torch.linalg.solve_triangular(L, K_nm, upper=False)
        var = K_mm.diagonal(dim1=-2, dim2=-1) - (v ** 2).sum(dim=-2)
        var = var.clamp(min=1e-6)

        return mean, var


input_dim = 4
num_train = 128
num_test = 32
batch_size = 8


def get_inputs():
    p = popcorn_pri
    X_train = torch.randn(p.jitter_int(batch_size), p.jitter_int(num_train), input_dim)
    Y_train = torch.randn(p.jitter_int(batch_size), p.jitter_int(num_train), 1)
    X_test = torch.randn(p.jitter_int(batch_size), p.jitter_int(num_test), input_dim)
    return [X_train, Y_train, X_test]


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>

__global__ void rbf_kernel_kernel(
    const float* __restrict__ X1,
    const float* __restrict__ X2,
    float* __restrict__ K,
    int B, int N1, int N2, int D,
    float l2, float var) {

    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * N1 * N2;
    if (idx >= total) return;

    int b = idx / (N1 * N2);
    int rem = idx % (N1 * N2);
    int i = rem / N2;
    int j = rem % N2;

    float dist_sq = 0.0f;
    for (int d = 0; d < D; d++) {
        float diff = X1[b * N1 * D + i * D + d] - X2[b * N2 * D + j * D + d];
        dist_sq += diff * diff;
    }
    K[idx] = var * expf(-0.5f * dist_sq / l2);
}

torch::Tensor rbf_kernel_cuda(torch::Tensor X1, torch::Tensor X2, float lengthscale, float variance) {
    auto B = X1.size(0);
    auto N1 = X1.size(1);
    auto N2 = X2.size(1);
    auto D = X1.size(2);

    auto K = torch::zeros({B, N1, N2}, X1.options());

    float l2 = lengthscale * lengthscale;
    float var = variance;

    int total = B * N1 * N2;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;

    rbf_kernel_kernel<<<blocks, threads>>>(
        X1.data_ptr<float>(), X2.data_ptr<float>(),
        K.data_ptr<float>(),
        B, N1, N2, D, l2, var);

    return K;
}
'''

cpp_source = '''
torch::Tensor rbf_kernel_cuda(torch::Tensor X1, torch::Tensor X2, float lengthscale, float variance);
'''

ext = load_inline(
    name='rbf_ext',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['rbf_kernel_cuda'],
    verbose=False
)

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

    def forward(self, X_train, Y_train, X_test):
        l = self.log_lengthscale.exp()
        v = self.log_variance.exp()

        K_nn = self.ext.rbf_kernel_cuda(X_train, X_train, l.item(), v.item())
        eye = torch.eye(K_nn.shape[-1], device=K_nn.device).unsqueeze(0)
        K_nn = K_nn + self.noise * eye
        K_nm = self.ext.rbf_kernel_cuda(X_train, X_test, l.item(), v.item())
        K_mm = self.ext.rbf_kernel_cuda(X_test, X_test, l.item(), v.item())

        L = torch.linalg.cholesky(K_nn)
        alpha = torch.cholesky_solve(Y_train, L)

        mean = K_nm.transpose(-1, -2) @ alpha

        v_tri = torch.linalg.solve_triangular(L, K_nm, upper=False)
        var = K_mm.diagonal(dim1=-2, dim2=-1) - (v_tri ** 2).sum(dim=-2)
        var = var.clamp(min=1e-6)

        return mean, var