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

_cpp_src = r'''
#include <torch/extension.h>
#include <vector>
std::vector<torch::Tensor> bayes_forward(torch::Tensor X, torch::Tensor y, torch::Tensor X_test,
                                         double prior_precision, double noise_precision);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#define DMAX 16
#define NMAX 256

__global__ void bayes_kernel(const float* __restrict__ X,
                             const float* __restrict__ y,
                             const float* __restrict__ Xt,
                             float* __restrict__ mean,
                             float* __restrict__ var,
                             int B, int N, int M,
                             float alpha, float beta) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    __shared__ float Xs[NMAX*DMAX];
    __shared__ float ys[NMAX];
    __shared__ float A[DMAX*DMAX];
    __shared__ float L[DMAX*DMAX];
    __shared__ float rhs[DMAX];
    __shared__ float mu[DMAX];
    const float* Xb = X + (size_t)b * N * DMAX;
    const float* yb = y + (size_t)b * N;
    const float* Xtb = Xt + (size_t)b * M * DMAX;

    for (int idx = tid; idx < N*DMAX; idx += blockDim.x) Xs[idx] = Xb[idx];
    if (tid < N) ys[tid] = yb[tid];
    if (tid < DMAX*DMAX) { L[tid] = 0.0f; A[tid] = 0.0f; }
    __syncthreads();

    if (tid < DMAX*DMAX) {
        int i = tid >> 4;
        int j = tid & 15;
        if (j <= i) {
            float s = 0.0f;
            #pragma unroll 4
            for (int n = 0; n < NMAX; ++n) s += Xs[n*DMAX + i] * Xs[n*DMAX + j];
            A[tid] = beta * s + ((i == j) ? alpha : 0.0f);
        }
    }
    if (tid < DMAX) {
        float s = 0.0f;
        #pragma unroll 4
        for (int n = 0; n < NMAX; ++n) s += Xs[n*DMAX + tid] * ys[n];
        rhs[tid] = beta * s;
        mu[tid] = 0.0f;
    }
    __syncthreads();

    if (tid == 0) {
        for (int i = 0; i < DMAX; ++i) {
            for (int j = 0; j <= i; ++j) {
                float s = A[i*DMAX + j];
                for (int k = 0; k < j; ++k) s -= L[i*DMAX + k] * L[j*DMAX + k];
                if (i == j) L[i*DMAX + j] = sqrtf(fmaxf(s, 0.0f));
                else L[i*DMAX + j] = s / L[j*DMAX + j];
            }
        }
        for (int i = 0; i < DMAX; ++i) {
            float s = rhs[i];
            for (int k = 0; k < i; ++k) s -= L[i*DMAX + k] * rhs[k];
            rhs[i] = s / L[i*DMAX + i];
        }
        for (int i = DMAX-1; i >= 0; --i) {
            float s = rhs[i];
            for (int k = i+1; k < DMAX; ++k) s -= L[k*DMAX + i] * mu[k];
            mu[i] = s / L[i*DMAX + i];
        }
    }
    __syncthreads();

    for (int m = tid; m < M; m += blockDim.x) {
        const float* x = Xtb + (size_t)m * DMAX;
        float pm = 0.0f;
        #pragma unroll
        for (int i = 0; i < DMAX; ++i) pm += x[i] * mu[i];
        float z[DMAX];
        float vv = 0.0f;
        #pragma unroll
        for (int i = 0; i < DMAX; ++i) {
            float s = x[i];
            #pragma unroll
            for (int k = 0; k < i; ++k) s -= L[i*DMAX + k] * z[k];
            float zi = s / L[i*DMAX + i];
            z[i] = zi;
            vv += zi * zi;
        }
        mean[(size_t)b * M + m] = pm;
        var[(size_t)b * M + m] = vv + 1.0f / beta;
    }
}

std::vector<torch::Tensor> bayes_forward(torch::Tensor X, torch::Tensor y, torch::Tensor X_test,
                                         double prior_precision, double noise_precision) {
    TORCH_CHECK(X.is_cuda() && y.is_cuda() && X_test.is_cuda(), "inputs must be CUDA tensors");
    TORCH_CHECK(X.scalar_type() == torch::kFloat32, "only float32 supported");
    TORCH_CHECK(X.size(2) == DMAX && X.size(1) == NMAX, "expected N=256,D=16");
    auto Xc = X.contiguous();
    auto yc = y.contiguous();
    auto Xtc = X_test.contiguous();
    int B = (int)Xc.size(0);
    int N = (int)Xc.size(1);
    int M = (int)Xtc.size(1);
    auto mean = torch::empty({B, M, 1}, X.options());
    auto var = torch::empty({B, M}, X.options());
    bayes_kernel<<<B, 256>>>(Xc.data_ptr<float>(), yc.data_ptr<float>(), Xtc.data_ptr<float>(),
                             mean.data_ptr<float>(), var.data_ptr<float>(), B, N, M,
                             (float)prior_precision, (float)noise_precision);
    return {mean, var};
}
'''

_ext = load_inline(name='bayes_lr_ext_final', cpp_sources=_cpp_src, cuda_sources=_cuda_src,
                   functions=['bayes_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'], 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: torch.Tensor, y: torch.Tensor, X_test: torch.Tensor) -> tuple:
        out = _ext.bayes_forward(X, y, X_test, float(self.prior_precision), float(self.noise_precision))
        return out[0], out[1]