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_runtime.h>
#include <vector>

__device__ void cholesky_decomp(float* A, float* L, int D) {
    for (int i = 0; i < D; i++) {
        for (int j = 0; j <= i; j++) {
            float sum = 0.0f;
            for (int k = 0; k < j; k++) {
                sum += L[i*D + k] * L[j*D + k];
            }
            if (i == j) {
                L[i*D + j] = sqrtf(A[i*D + i] - sum);
            } else {
                L[i*D + j] = (A[i*D + j] - sum) / L[j*D + j];
            }
        }
    }
}

__device__ void forward_sub(const float* L, const float* b, float* y, int D) {
    for (int i = 0; i < D; i++) {
        float sum = 0.0f;
        for (int j = 0; j < i; j++) sum += L[i*D + j] * y[j];
        y[i] = (b[i] - sum) / L[i*D + i];
    }
}

__device__ void back_sub(const float* L, const float* y, float* x, int D) {
    for (int i = D-1; i >= 0; i--) {
        float sum = 0.0f;
        for (int j = i+1; j < D; j++) sum += L[j*D + i] * x[j];
        x[i] = (y[i] - sum) / L[i*D + i];
    }
}

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

    extern __shared__ float sm[];

    float* Lambda = sm;
    float* L = &sm[D*D];
    float* invL = &sm[2*D*D];
    float* rhs = &sm[3*D*D];
    float* mu = &sm[3*D*D + D];
    float* col = &sm[3*D*D + 2*D];
    float* tmpvec = &sm[3*D*D + 3*D];

    // X^T X
    for (int i = 0; i < D; i++) {
        for (int j = 0; j < D; j++) {
            float s = 0.0f;
            for (int k = 0; k < N; k++) {
                s += X[(b*N + k)*D + i] * X[(b*N + k)*D + j];
            }
            Lambda[i*D + j] = beta * s;
        }
    }
    for (int i = 0; i < D; i++) Lambda[i*D + i] += alpha;

    cholesky_decomp(Lambda, L, D);

    // rhs
    for (int i = 0; i < D; i++) {
        float s = 0.0f;
        for (int k = 0; k < N; k++) {
            s += X[(b*N + k)*D + i] * y[b*N + k];
        }
        rhs[i] = beta * s;
    }

    forward_sub(L, rhs, col, D);
    back_sub(L, col, mu, D);

    // pred_mean
    for (int m = 0; m < M; m++) {
        float s = 0.0f;
        for (int d = 0; d < D; d++) {
            s += X_test[(b*M + m)*D + d] * mu[d];
        }
        pred_mean[b*M + m] = s;
    }

    // inv(L) columns , note L lower, solve gives inv(L)
    for (int c = 0; c < D; c++) {
        float e[32] = {0.0f};
        e[c] = 1.0f;
        forward_sub(L, e, col, D);
        for (int r=0; r<D; r++) invL[r*D + c] = col[r];
    }

    // pred_var: || Xrow @ invL || ^2 + 1/beta
    for (int m = 0; m < M; m++) {
        for (int d=0; d<D; d++) {
            float t = 0.0f;
            for (int k=0; k<D; k++) t += X_test[(b*M + m)*D + k] * invL[k*D + d];
            tmpvec[d] = t;
        }
        float s = 0.0f;
        for (int d=0; d<D; d++) s += tmpvec[d] * tmpvec[d];
        pred_var[b*M + m] = s + 1.0f / beta;
    }
}

std::vector<torch::Tensor> bayes_forward(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().device(X.device()).dtype(X.dtype());
    auto pred_mean = torch::empty({B, M}, options);
    auto pred_var = torch::empty({B, M}, options);
    size_t smem = (4*D*D + 4*D) * sizeof(float);
    compute_bayes<<<B, 1, smem>>>(
        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
    );
    return {pred_mean.view({B, M, 1}), pred_var};
}
'''

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

module = load_inline(
    name='bayes_linear',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['bayes_forward'],
    verbose=True,
    extra_cuda_cflags=['-O3']
)

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
        self.module = module

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