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>
#include <cmath>

constexpr int TILE_SIZE = 32;

template <typename scalar_t>
__global__ void rbf_kernel_kernel(
    const scalar_t* __restrict__ X1,
    const scalar_t* __restrict__ X2,
    const scalar_t* __restrict__ X1_norm_sq,
    const scalar_t* __restrict__ X2_norm_sq,
    scalar_t* __restrict__ output,
    const scalar_t lengthscale_sq,
    const scalar_t variance,
    int batch_size,
    int N,
    int M,
    int D) {
    
    const int batch_idx = blockIdx.z;
    const int row = blockIdx.y * blockDim.y + threadIdx.y;
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (batch_idx >= batch_size || row >= N || col >= M) return;
    
    // Compute dot product with unrolled loop
    scalar_t dot = 0;
    const scalar_t* x1_ptr = X1 + batch_idx * N * D + row * D;
    const scalar_t* x2_ptr = X2 + batch_idx * M * D + col * D;
    
    for (int d = 0; d < D; d++) {
        dot += x1_ptr[d] * x2_ptr[d];
    }
    
    // dist_sq = ||X1[row]||^2 + ||X2[col]||^2 - 2 * dot
    scalar_t dist_sq = X1_norm_sq[batch_idx * N + row] + 
                       X2_norm_sq[batch_idx * M + col] - 2.0f * dot;
    
    // RBF kernel: variance * exp(-0.5 * dist_sq / lengthscale_sq)
    output[batch_idx * N * M + row * M + col] = variance * expf(-0.5f * dist_sq / lengthscale_sq);
}

// Kernel for K_nn with noise diagonal already added
template <typename scalar_t>
__global__ void rbf_kernel_knn_kernel(
    const scalar_t* __restrict__ X1,
    const scalar_t* __restrict__ X1_norm_sq,
    scalar_t* __restrict__ output,
    const scalar_t lengthscale_sq,
    const scalar_t variance,
    const scalar_t noise,
    int batch_size,
    int N,
    int D) {
    
    const int batch_idx = blockIdx.z;
    const int row = blockIdx.y * blockDim.y + threadIdx.y;
    const int col = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (batch_idx >= batch_size || row >= N || col >= N) return;
    
    // Compute dot product
    scalar_t dot = 0;
    const scalar_t* x1_ptr = X1 + batch_idx * N * D + row * D;
    const scalar_t* x2_ptr = X1 + batch_idx * N * D + col * D;
    
    for (int d = 0; d < D; d++) {
        dot += x1_ptr[d] * x2_ptr[d];
    }
    
    // dist_sq = ||X1[row]||^2 + ||X1[col]||^2 - 2 * dot
    scalar_t dist_sq = X1_norm_sq[batch_idx * N + row] + 
                       X1_norm_sq[batch_idx * N + col] - 2.0f * dot;
    
    // RBF kernel: variance * exp(-0.5 * dist_sq / lengthscale_sq)
    scalar_t val = variance * expf(-0.5f * dist_sq / lengthscale_sq);
    
    // Add noise on diagonal
    if (row == col) {
        val += noise;
    }
    
    output[batch_idx * N * N + row * N + col] = val;
}

torch::Tensor rbf_kernel_cuda(
    torch::Tensor X1, 
    torch::Tensor X2, 
    double lengthscale, 
    double variance) {
    
    const int batch_size = X1.size(0);
    const int N = X1.size(1);
    const int M = X2.size(1);
    const int D = X1.size(2);
    
    // Compute squared norms
    auto X1_norm_sq = (X1 * X1).sum(-1);  // (B, N)
    auto X2_norm_sq = (X2 * X2).sum(-1);  // (B, M)
    
    auto output = torch::empty({batch_size, N, M}, X1.options());
    
    const dim3 threads(16, 16);
    const dim3 blocks((M + 15) / 16, (N + 15) / 16, batch_size);
    
    AT_DISPATCH_FLOATING_TYPES(X1.scalar_type(), "rbf_kernel_kernel", ([&] {
        rbf_kernel_kernel<scalar_t><<<blocks, threads>>>(
            X1.data_ptr<scalar_t>(),
            X2.data_ptr<scalar_t>(),
            X1_norm_sq.data_ptr<scalar_t>(),
            X2_norm_sq.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            lengthscale * lengthscale,
            variance,
            batch_size, N, M, D);
    }));
    
    return output;
}

torch::Tensor rbf_kernel_knn_cuda(
    torch::Tensor X1, 
    double lengthscale, 
    double variance,
    double noise) {
    
    const int batch_size = X1.size(0);
    const int N = X1.size(1);
    const int D = X1.size(2);
    
    // Compute squared norms
    auto X1_norm_sq = (X1 * X1).sum(-1);  // (B, N)
    
    auto output = torch::empty({batch_size, N, N}, X1.options());
    
    const dim3 threads(16, 16);
    const dim3 blocks((N + 15) / 16, (N + 15) / 16, batch_size);
    
    AT_DISPATCH_FLOATING_TYPES(X1.scalar_type(), "rbf_kernel_knn_kernel", ([&] {
        rbf_kernel_knn_kernel<scalar_t><<<blocks, threads>>>(
            X1.data_ptr<scalar_t>(),
            X1_norm_sq.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            lengthscale * lengthscale,
            variance,
            noise,
            batch_size, N, D);
    }));
    
    return output;
}
"""

cpp_source = """
torch::Tensor rbf_kernel_cuda(torch::Tensor X1, torch::Tensor X2, double lengthscale, double variance);
torch::Tensor rbf_kernel_knn_cuda(torch::Tensor X1, double lengthscale, double variance, double noise);
"""

gp_module = load_inline(
    name='gp_kernel_cuda_v4',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['rbf_kernel_cuda', 'rbf_kernel_knn_cuda'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
)

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())
    
    def forward(self, X_train, Y_train, X_test):
        l = self.log_lengthscale.exp().item()
        v = self.log_variance.exp().item()
        
        # Compute K_nn with noise diagonal fused
        K_nn = gp_module.rbf_kernel_knn_cuda(X_train, l, v, self.noise)
        
        # Compute K_nm (train x test)
        K_nm = gp_module.rbf_kernel_cuda(X_train, X_test, l, v)
        
        # K_mm diagonal is just variance (dist_sq = 0 for self-pairs)
        batch_size, M = X_test.shape[0], X_test.shape[1]
        K_mm_diag = torch.full((batch_size, M), v, dtype=X_test.dtype, device=X_test.device)
        
        # Cholesky decomposition and solve
        L = torch.linalg.cholesky(K_nn)
        alpha = torch.cholesky_solve(Y_train, L)
        
        # Predictive mean: K_nm.T @ alpha = K_mn @ alpha
        mean = K_nm.transpose(-1, -2) @ alpha
        
        # Predictive variance
        v_matrix = torch.linalg.solve_triangular(L, K_nm, upper=False)
        var = K_mm_diag - (v_matrix ** 2).sum(dim=-2)
        var = var.clamp(min=1e-6)
        
        return mean, var