import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Bayesian linear regression with a conjugate Gaussian prior.
    Given features X and targets y, computes the posterior distribution
    over weights: p(w | X, y) = N(mu_post, Sigma_post) in closed form
    via the matrix inversion lemma / Cholesky decomposition.
    """

    def __init__(self, input_dim, prior_precision=1.0, noise_precision=1.0):
        super().__init__()
        self.input_dim = input_dim
        self.prior_precision = prior_precision
        self.noise_precision = noise_precision

    def forward(
        self, X: torch.Tensor, y: torch.Tensor, X_test: torch.Tensor
    ) -> tuple:
        """
        Args:
            X:      (B, N, D) – training features
            y:      (B, N, 1) – training targets
            X_test: (B, M, D) – test features
        Returns:
            (pred_mean, pred_var):
                pred_mean: (B, M, 1)
                pred_var:  (B, M)
        """
        D = X.shape[-1]
        eye = torch.eye(D, device=X.device, dtype=X.dtype).unsqueeze(0)

        # Posterior precision: Lambda_post = alpha * I + beta * X^T X
        Lambda_post = self.prior_precision * eye + self.noise_precision * (X.transpose(-1, -2) @ X)
        L = torch.linalg.cholesky(Lambda_post)

        # Posterior mean: mu_post = Lambda_post^{-1} * beta * X^T y
        rhs = self.noise_precision * (X.transpose(-1, -2) @ y)
        mu_post = torch.cholesky_solve(rhs, L)

        # Predictive distribution
        pred_mean = X_test @ mu_post
        Sigma_post = torch.cholesky_inverse(L)
        pred_var = (X_test @ Sigma_post * X_test).sum(-1) + 1.0 / self.noise_precision

        return pred_mean, pred_var


input_dim = 16
num_train = 256
num_test = 64
batch_size = 16


def get_inputs():
    p = popcorn_pri
    X = torch.randn(p.jitter_int(batch_size), p.jitter_int(num_train), input_dim)
    y = 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, y, 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 = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void bayesian_linear_kernel(
    const float* __restrict__ X,
    const float* __restrict__ y,
    const float* __restrict__ X_test,
    float* __restrict__ pred_mean,
    float* __restrict__ pred_var,
    int B, int N, int D, int M,
    float prior_precision, float noise_precision
) {
    int b = blockIdx.x;
    if (b >= B) return;

    extern __shared__ float smem[];

    const float* Xb = X + b * N * D;
    const float* yb = y + b * N;
    const float* Xt = X_test + b * M * D;

    float* Lambda = smem;
    float* L = smem + D*D;
    float* rhs = smem + 2*D*D;
    float* mu_post = smem + 2*D*D + D;
    float* Sigma = smem + 2*D*D + 2*D;
    float* temp = smem + 3*D*D + 2*D;

    // Initialize Lambda = alpha * I
    for (int i = threadIdx.x; i < D*D; i += blockDim.x) {
        int r = i / D, c = i % D;
        Lambda[i] = (r == c) ? prior_precision : 0.0f;
    }
    __syncthreads();

    // Compute beta * X^T X and add to Lambda, X^T y
    for (int i = threadIdx.x; i < D; i += blockDim.x) {
        for (int j = 0; j < D; j++) {
            float sum_xx = 0.0f;
            for (int k = 0; k < N; k++) {
                sum_xx += Xb[k*D + i] * Xb[k*D + j];
            }
            Lambda[i*D + j] += noise_precision * sum_xx;
        }
        float sum_xy = 0.0f;
        for (int k = 0; k < N; k++) {
            sum_xy += Xb[k*D + i] * yb[k];
        }
        rhs[i] = noise_precision * sum_xy;
    }
    __syncthreads();

    // Cholesky L of Lambda, lower triangular
    for (int i = 0; i < D; ++i) {
        for (int j = threadIdx.x; j <= i; j += blockDim.x) {
            float s = 0.0f;
            for (int k = 0; k < j; ++k) {
                s += L[i*D + k] * L[j*D + k];
            }
            if (i == j) {
                L[i*D + i] = sqrtf(Lambda[i*D + i] - s);
            } else {
                L[i*D + j] = (Lambda[i*D + j] - s) / L[j*D + j];
            }
        }
        __syncthreads();
    }

    // Forward substitution L z = rhs -> mu_post temporarily
    for (int i = 0; i < D; ++i) {
        if (threadIdx.x == 0) {
            float s = 0.0f;
            for (int k = 0; k < i; ++k) s += L[i*D + k] * mu_post[k];
            mu_post[i] = (rhs[i] - s) / L[i*D + i];
        }
        __syncthreads();
    }
    // Backward L^T mu = z
    for (int i = D - 1; i >= 0; --i) {
        if (threadIdx.x == 0) {
            float s = 0.0f;
            for (int k = i + 1; k < D; ++k) s += L[k*D + i] * mu_post[k];
            mu_post[i] = (mu_post[i] - s) / L[i*D + i];
        }
        __syncthreads();
    }

    // pred_mean = X_test @ mu_post
    for (int m = threadIdx.x; m < M; m += blockDim.x) {
        float sum = 0.0f;
        for (int d = 0; d < D; ++d) {
            sum += Xt[m * D + d] * mu_post[d];
        }
        pred_mean[b * M + m] = sum;
    }

    // Compute Sigma = inv(Lambda) via cholesky inverse: solve L Y = I then L^T Sigma = Y
    for (int col = 0; col < D; ++col) {
        // solve L Y_col = e_col
        for (int i = 0; i < D; ++i) {
            if (threadIdx.x == 0) {
                float val = (i == col ? 1.0f : 0.0f);
                for (int k = 0; k < i; ++k) val -= L[i*D + k] * temp[k*D + col];
                temp[i*D + col] = val / L[i*D + i];
            }
            __syncthreads();
        }
        // solve L^T Sigma_col = Y_col
        for (int i = D-1; i >= 0; --i) {
            if (threadIdx.x == 0) {
                float val = temp[i*D + col];
                for (int k = i+1; k < D; ++k) val -= L[k*D + i] * Sigma[k*D + col];
                Sigma[i*D + col] = val / L[i*D + i];
            }
            __syncthreads();
        }
    }

    // pred_var = diag(X_test @ Sigma @ X_test^T) + 1/beta
    for (int m = threadIdx.x; m < M; m += blockDim.x) {
        float sum = 0.0f;
        for (int d = 0; d < D; ++d) {
            float tmp = 0.0f;
            for (int e = 0; e < D; ++e) {
                tmp += Xt[m*D + e] * Sigma[e * D + d];
            }
            sum += tmp * Xt[m*D + d];
        }
        pred_var[b * M + m] = sum + 1.0f / noise_precision;
    }
}

std::vector<torch::Tensor> bayesian_linear_cuda(
    torch::Tensor X, torch::Tensor y, torch::Tensor X_test,
    float prior_precision, float noise_precision
) {
    int B = X.size(0);
    int N = X.size(1);
    int D = X.size(2);
    int M = X_test.size(1);

    auto options = torch::TensorOptions().dtype(X.dtype()).device(X.device());
    auto pred_mean = torch::zeros({B, M, 1}, options);
    auto pred_var = torch::zeros({B, M}, options);

    dim3 grid(B);
    dim3 block(256);
    size_t smem_size = (4 * D * D + 2 * D) * sizeof(float);

    bayesian_linear_kernel<<<grid, block, smem_size>>>(
        X.data_ptr<float>(), y.data_ptr<float>(), X_test.data_ptr<float>(),
        pred_mean.data_ptr<float>(), pred_var.data_ptr<float>(),
        B, N, D, M, prior_precision, noise_precision
    );

    cudaDeviceSynchronize();
    return {pred_mean.view({B, M, 1}), pred_var};
}
'''

cpp_source = '''
std::vector<torch::Tensor> bayesian_linear_cuda(
    torch::Tensor X, torch::Tensor y, torch::Tensor X_test,
    float prior_precision, float noise_precision
);
'''

module = load_inline(
    name='bayesian_linear',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['bayesian_linear_cuda'],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self, input_dim, prior_precision=1.0, noise_precision=1.0):
        super().__init__()
        self.input_dim = input_dim
        self.prior_precision = prior_precision
        self.noise_precision = noise_precision

    def forward(self, X, y, X_test):
        res = module.bayesian_linear_cuda(X, y, X_test, self.prior_precision, self.noise_precision)
        return res[0], res[1]