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

std::vector<torch::Tensor> gp_forward_cuda(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.h>
#include <cuda_runtime.h>
#include <vector>
#include <math.h>

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

__device__ __forceinline__ int lidx(int i, int j) {
    return (i * (i + 1)) / 2 + j;
}

__device__ __forceinline__ float warp_sum(float v) {
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xffffffff, v, off);
    return v;
}

__global__ void gp_kernel(const float* __restrict__ Xtr,
                          const float* __restrict__ Y,
                          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, int N, int D, int M) {
    int b = blockIdx.x;
    if (b >= B) return;
    int tid = threadIdx.x;

    extern __shared__ float sh[];
    int pack = N * (N + 1) / 2;
    float* L = sh;                         // packed lower triangular matrix
    float* alpha = sh + pack;              // length N, first z then alpha
    float* vbuf = alpha + N;               // 4 x N triangular solve workspace

    float ell = expf(log_l_ptr[0]);
    float inv_2_l2 = 0.5f / (ell * ell);
    float sig2 = expf(log_v_ptr[0]);

    const float* Xb = Xtr + ((long long)b * N * D);
    const float* Yb = Y + ((long long)b * N);
    const float* Tst = Xte + ((long long)b * M * D);

    // Build the lower half of K(X_train, X_train) + noise*I.
    int NN = N * N;
    for (int t = tid; t < NN; t += blockDim.x) {
        int i = t / N;
        int j = t - i * N;
        if (i >= j) {
            float ds = 0.0f;
            for (int d = 0; d < D; ++d) {
                float diff = Xb[i * D + d] - Xb[j * D + d];
                ds += diff * diff;
            }
            float val = sig2 * expf(-ds * inv_2_l2);
            if (i == j) val += noise;
            L[lidx(i, j)] = val;
        }
    }
    __syncthreads();

    // In-place Cholesky factorization of the packed matrix.  One block handles
    // one GP, with parallelism across rows for each column.
    for (int j = 0; j < N; ++j) {
        if (tid == 0) {
            float s = L[lidx(j, j)];
            for (int k = 0; k < j; ++k) {
                float ljk = L[lidx(j, k)];
                s -= ljk * ljk;
            }
            L[lidx(j, j)] = sqrtf(fmaxf(s, 1.0e-20f));
        }
        __syncthreads();
        float diag = L[lidx(j, j)];
        for (int i = j + 1 + tid; i < N; i += blockDim.x) {
            float s = L[lidx(i, j)];
            for (int k = 0; k < j; ++k) {
                s -= L[lidx(i, k)] * L[lidx(j, k)];
            }
            L[lidx(i, j)] = s / diag;
        }
        __syncthreads();
    }

    // Solve K alpha = y through the Cholesky factor.  Single RHS, so the
    // sequential solve is tiny compared with the factorization.
    if (tid == 0) {
        for (int i = 0; i < N; ++i) {
            float s = Yb[i];
            for (int k = 0; k < i; ++k) s -= L[lidx(i, k)] * alpha[k];
            alpha[i] = s / L[lidx(i, i)];
        }
        for (int i = N - 1; i >= 0; --i) {
            float s = alpha[i];
            for (int k = i + 1; k < N; ++k) s -= L[lidx(k, i)] * alpha[k];
            alpha[i] = s / L[lidx(i, i)];
        }
    }
    __syncthreads();

    // Predictive mean and variance.  Process four test points at a time, one
    // warp per test point.  The forward triangular solve dot products are
    // parallelized within the warp.
    int warp = tid >> 5;
    int lane = tid & 31;
    for (int g = 0; g < M; g += 4) {
        int m = g + warp;
        float mean_acc = 0.0f;
        float vv_acc = 0.0f;
        if (m < M) {
            float* vw = vbuf + warp * N;
            const float* Xt = Tst + m * D;
            for (int i = 0; i < N; ++i) {
                float kval = 0.0f;
                if (lane == 0) {
                    float ds = 0.0f;
                    for (int d = 0; d < D; ++d) {
                        float diff = Xb[i * D + d] - Xt[d];
                        ds += diff * diff;
                    }
                    kval = sig2 * expf(-ds * inv_2_l2);
                }
                float dot = 0.0f;
                for (int k = lane; k < i; k += 32) {
                    dot += L[lidx(i, k)] * vw[k];
                }
                dot = warp_sum(dot);
                if (lane == 0) {
                    float vi = (kval - dot) / L[lidx(i, i)];
                    vw[i] = vi;
                    mean_acc += kval * alpha[i];
                    vv_acc += vi * vi;
                }
                __syncwarp();
            }
            if (lane == 0) {
                float vr = sig2 - vv_acc;
                if (vr < 1.0e-6f) vr = 1.0e-6f;
                mean[((long long)b * M + m)] = mean_acc;
                var_out[(long long)b * M + m] = vr;
            }
        }
        __syncthreads();
    }
}

std::vector<torch::Tensor> gp_forward_cuda(torch::Tensor X_train,
                                           torch::Tensor Y_train,
                                           torch::Tensor X_test,
                                           torch::Tensor log_lengthscale,
                                           torch::Tensor log_variance,
                                           double noise) {
    CHECK_INPUT(X_train);
    CHECK_INPUT(Y_train);
    CHECK_INPUT(X_test);
    CHECK_INPUT(log_lengthscale);
    CHECK_INPUT(log_variance);
    TORCH_CHECK(X_train.scalar_type() == torch::kFloat32, "X_train must be float32");
    TORCH_CHECK(Y_train.scalar_type() == torch::kFloat32, "Y_train must be float32");
    TORCH_CHECK(X_test.scalar_type() == torch::kFloat32, "X_test must be float32");

    int B = (int)X_train.size(0);
    int N = (int)X_train.size(1);
    int D = (int)X_train.size(2);
    int M = (int)X_test.size(1);

    auto mean = torch::empty({B, M, 1}, X_train.options());
    auto var = torch::empty({B, M}, X_train.options());

    int threads = 128;
    int pack = N * (N + 1) / 2;
    size_t shmem = (size_t)(pack + N + 4 * N) * sizeof(float);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    gp_kernel<<<B, threads, shmem, stream>>>(
        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>(), B, N, D, M);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return {mean, var};
}
"""

_gp_ext = load_inline(
    name="gp_cholesky_predict_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["gp_forward_cuda"],
    extra_cuda_cflags=["-O3"],
    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(_gp_ext.gp_forward_cuda(X_train, Y_train, X_test, self.log_lengthscale, self.log_variance, float(self.noise)))