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

cpp_src = r'''
#include <torch/extension.h>
std::vector<torch::Tensor> gp_forward(torch::Tensor X_train, torch::Tensor Y_train, torch::Tensor X_test,
                                      torch::Tensor log_lengthscale, torch::Tensor log_variance, double noise);
'''

cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <vector>
#define NVAL 128
#define MVAL 32
#define DVAL 4

__inline__ __device__ float warp_sum(float v) {
    v += __shfl_down_sync(0xffffffff, v, 16);
    v += __shfl_down_sync(0xffffffff, v, 8);
    v += __shfl_down_sync(0xffffffff, v, 4);
    v += __shfl_down_sync(0xffffffff, v, 2);
    v += __shfl_down_sync(0xffffffff, v, 1);
    return v;
}

__global__ void gp_forward_kernel(const float* __restrict__ Xtr,
                                  const float* __restrict__ Ytr,
                                  const float* __restrict__ Xte,
                                  const float* __restrict__ log_l_ptr,
                                  const float* __restrict__ log_v_ptr,
                                  float noise,
                                  float* __restrict__ mean,
                                  float* __restrict__ var_out) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    int warp = tid >> 5;
    int lane = tid & 31;
    extern __shared__ float sh[];
    float* K = sh;
    float* alpha = K + NVAL*NVAL;
    float* V = alpha + NVAL;

    float l2 = expf(2.0f * (*log_l_ptr));
    float coef = -0.5f / l2;
    float variance = expf(*log_v_ptr);

    const float* Xt = Xtr + b * NVAL * DVAL;
    const float* Y = Ytr + b * NVAL;
    const float* Xs = Xte + b * MVAL * DVAL;

    for (int idx = tid; idx < NVAL * NVAL; idx += blockDim.x) {
        int r = idx >> 7;
        int c = idx & 127;
        float val = 0.0f;
        if (r >= c) {
            float dx0 = Xt[(r<<2)]   - Xt[(c<<2)];
            float dx1 = Xt[(r<<2)+1] - Xt[(c<<2)+1];
            float dx2 = Xt[(r<<2)+2] - Xt[(c<<2)+2];
            float dx3 = Xt[(r<<2)+3] - Xt[(c<<2)+3];
            float dist = dx0*dx0 + dx1*dx1 + dx2*dx2 + dx3*dx3;
            val = variance * __expf(coef * dist);
            if (r == c) val += noise;
        }
        K[idx] = val;
    }
    __syncthreads();

    for (int k = 0; k < NVAL; ++k) {
        if (tid == 0) {
            float diagv = K[k * NVAL + k];
            diagv = diagv > 1.0e-20f ? diagv : 1.0e-20f;
            K[k * NVAL + k] = sqrtf(diagv);
        }
        __syncthreads();
        float diag = K[k * NVAL + k];
        for (int i = k + 1 + tid; i < NVAL; i += blockDim.x) {
            K[i * NVAL + k] = K[i * NVAL + k] / diag;
        }
        __syncthreads();
        for (int i = k + 1 + warp; i < NVAL; i += 8) {
            float lik = K[i * NVAL + k];
            for (int j = k + 1 + lane; j <= i; j += 32) {
                K[i * NVAL + j] -= lik * K[j * NVAL + k];
            }
        }
        __syncthreads();
    }

    if (warp == 0) {
        for (int i = 0; i < NVAL; ++i) {
            float dot = 0.0f;
            for (int j = lane; j < i; j += 32) dot += K[i * NVAL + j] * alpha[j];
            dot = warp_sum(dot);
            if (lane == 0) alpha[i] = (Y[i] - dot) / K[i * NVAL + i];
            __syncwarp();
        }
        for (int i = NVAL - 1; i >= 0; --i) {
            float dot = 0.0f;
            for (int j = i + 1 + lane; j < NVAL; j += 32) dot += K[j * NVAL + i] * alpha[j];
            dot = warp_sum(dot);
            if (lane == 0) alpha[i] = (alpha[i] - dot) / K[i * NVAL + i];
            __syncwarp();
        }
    }
    __syncthreads();

    for (int idx = tid; idx < NVAL * MVAL; idx += blockDim.x) {
        int i = idx >> 5;
        int m = idx & 31;
        float dx0 = Xt[(i<<2)]   - Xs[(m<<2)];
        float dx1 = Xt[(i<<2)+1] - Xs[(m<<2)+1];
        float dx2 = Xt[(i<<2)+2] - Xs[(m<<2)+2];
        float dx3 = Xt[(i<<2)+3] - Xs[(m<<2)+3];
        float dist = dx0*dx0 + dx1*dx1 + dx2*dx2 + dx3*dx3;
        V[idx] = variance * __expf(coef * dist);
    }
    __syncthreads();

    for (int g = 0; g < 4; ++g) {
        int m = (g << 3) + warp;
        float s = 0.0f;
        for (int i = lane; i < NVAL; i += 32) s += V[i * MVAL + m] * alpha[i];
        s = warp_sum(s);
        if (lane == 0) mean[b * MVAL + m] = s;
    }
    __syncthreads();

    for (int i = 0; i < NVAL; ++i) {
        float diag = K[i * NVAL + i];
        for (int g = 0; g < 4; ++g) {
            int m = (g << 3) + warp;
            float dot = 0.0f;
            for (int j = lane; j < i; j += 32) dot += K[i * NVAL + j] * V[j * MVAL + m];
            dot = warp_sum(dot);
            if (lane == 0) V[i * MVAL + m] = (V[i * MVAL + m] - dot) / diag;
        }
        __syncthreads();
    }

    for (int g = 0; g < 4; ++g) {
        int m = (g << 3) + warp;
        float ss = 0.0f;
        for (int i = lane; i < NVAL; i += 32) {
            float x = V[i * MVAL + m];
            ss += x * x;
        }
        ss = warp_sum(ss);
        if (lane == 0) {
            float vv = variance - ss;
            if (vv < 1.0e-6f) vv = 1.0e-6f;
            var_out[b * MVAL + m] = vv;
        }
    }
}

std::vector<torch::Tensor> gp_forward(torch::Tensor X_train, torch::Tensor Y_train, torch::Tensor X_test,
                                      torch::Tensor log_lengthscale, torch::Tensor log_variance, double noise) {
    int B = X_train.size(0);
    auto mean = torch::empty({B, MVAL, 1}, X_train.options());
    auto var = torch::empty({B, MVAL}, X_train.options());
    int threads = 256;
    size_t shmem = (size_t)(NVAL * NVAL + NVAL + NVAL * MVAL) * sizeof(float);
    cudaFuncSetAttribute((const void*)gp_forward_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)shmem);
    gp_forward_kernel<<<B, threads, shmem, at::cuda::getCurrentCUDAStream()>>>(
        X_train.data_ptr<float>(), Y_train.data_ptr<float>(), X_test.data_ptr<float>(),
        log_lengthscale.data_ptr<float>(), log_variance.data_ptr<float>(), (float)noise,
        mean.data_ptr<float>(), var.data_ptr<float>());
    return {mean, var};
}
'''

_ext = load_inline(name='gp_custom_ext_v6', cpp_sources=cpp_src, cuda_sources=cuda_src, functions=['gp_forward'], with_cuda=True, extra_cuda_cflags=['-O3', '--use_fast_math'], extra_cflags=['-O3'], 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())
    def forward(self, X_train: torch.Tensor, Y_train: torch.Tensor, X_test: torch.Tensor) -> tuple:
        return tuple(_ext.gp_forward(X_train, Y_train, X_test, self.log_lengthscale, self.log_variance, float(self.noise)))