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 = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

constexpr int D = 16;

// Naive Cholesky decomposition for small matrices
__device__ void cholesky_decomp(const float* A, float* L) {
    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) {
                float val = A[i * D + i] - sum;
                L[i * D + j] = sqrtf(fmaxf(val, 1e-12f));
            } else {
                L[i * D + j] = (A[i * D + j] - sum) / L[j * D + j];
            }
        }
    }
}

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

// Backward substitution
__device__ void backward_solve(const float* L, const float* b, float* x) {
    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] = (b[i] - sum) / L[i * D + i];
    }
}

// Cholesky solve
__device__ void cholesky_solve(const float* L, const float* b, float* x, float* temp) {
    forward_solve(L, b, temp);
    backward_solve(L, temp, x);
}

// Cholesky inverse
__device__ void cholesky_inverse(const float* L, float* A_inv, float* temp1, float* temp2) {
    for (int col = 0; col < D; col++) {
        float e[16] = {0};
        e[col] = 1.0f;
        
        forward_solve(L, e, temp1);
        backward_solve(L, temp1, temp2);
        
        for (int row = 0; row < D; row++) {
            A_inv[row * D + col] = temp2[row];
        }
    }
}

// Main kernel with optimized memory access
__global__ void bayesian_linear_kernel(
    const float* __restrict__ X,      // (B, N, D)
    const float* __restrict__ y,      // (B, N, 1)
    const float* __restrict__ X_test, // (B, M, D)
    float* __restrict__ pred_mean,    // (B, M, 1)
    float* __restrict__ pred_var,     // (B, M)
    int B, int N, int M,
    float prior_precision, float noise_precision
) {
    int batch_idx = blockIdx.x;
    if (batch_idx >= B) return;
    
    int tid = threadIdx.x;
    int tcount = blockDim.x;
    
    // Shared memory - all in one block for efficiency
    __shared__ float smem[1024];
    float* XtX = smem;           // 256 floats
    float* L_mat = smem + 256;   // 256 floats
    float* Xty = smem + 512;     // 16 floats
    float* mu_post = smem + 528; // 16 floats
    float* Sigma_post = smem + 544; // 256 floats
    float* temp1 = smem + 800;   // 16 floats
    float* temp2 = smem + 816;   // 16 floats
    
    // Initialize XtX with prior: alpha * I
    for (int idx = tid; idx < D * D; idx += tcount) {
        int row = idx >> 4;
        int col = idx & 15;
        XtX[idx] = (row == col) ? prior_precision : 0.0f;
    }
    __syncthreads();
    
    // Compute X^T @ X: coalesced row-by-row loading
    __shared__ float x_row[16];
    for (int n = 0; n < N; n++) {
        if (tid < D) {
            x_row[tid] = X[batch_idx * N * D + n * D + tid];
        }
        __syncthreads();
        
        for (int idx = tid; idx < D * D; idx += tcount) {
            int i = idx >> 4;
            int j = idx & 15;
            XtX[idx] += noise_precision * x_row[i] * x_row[j];
        }
        __syncthreads();
    }
    
    // Cholesky decomposition (single thread)
    if (tid == 0) {
        cholesky_decomp(XtX, L_mat);
    }
    __syncthreads();
    
    // Compute X^T @ y
    for (int d = tid; d < D; d += tcount) {
        float sum = 0.0f;
        for (int n = 0; n < N; n++) {
            sum += X[batch_idx * N * D + n * D + d] * y[batch_idx * N + n];
        }
        Xty[d] = noise_precision * sum;
    }
    __syncthreads();
    
    // Solve for posterior mean
    if (tid == 0) {
        cholesky_solve(L_mat, Xty, mu_post, temp1);
    }
    __syncthreads();
    
    // Compute pred_mean = X_test @ mu_post
    for (int m = tid; m < M; m += tcount) {
        float sum = 0.0f;
        for (int d = 0; d < D; d++) {
            sum += X_test[batch_idx * M * D + m * D + d] * mu_post[d];
        }
        pred_mean[batch_idx * M + m] = sum;
    }
    
    // Compute Sigma_post = L^{-1} @ L^{-T}
    if (tid == 0) {
        cholesky_inverse(L_mat, Sigma_post, temp1, temp2);
    }
    __syncthreads();
    
    // Compute pred_var
    for (int m = tid; m < M; m += tcount) {
        float var = 0.0f;
        for (int i = 0; i < D; i++) {
            float sum = 0.0f;
            for (int j = 0; j < D; j++) {
                sum += Sigma_post[i * D + j] * X_test[batch_idx * M * D + m * D + j];
            }
            var += X_test[batch_idx * M * D + m * D + i] * sum;
        }
        pred_var[batch_idx * M + m] = var + 1.0f / noise_precision;
    }
}

std::tuple<torch::Tensor, torch::Tensor> bayesian_linear_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 M = X_test.size(1);
    
    auto pred_mean = torch::empty({B, M, 1}, X.options());
    auto pred_var = torch::empty({B, M}, X.options());
    
    int threads = 256;
    int blocks = B;
    
    bayesian_linear_kernel<<<blocks, threads>>>(
        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, M,
        prior_precision, noise_precision
    );
    
    return std::make_tuple(pred_mean, pred_var);
}
"""

cpp_source = """
std::tuple<torch::Tensor, torch::Tensor> bayesian_linear_forward(
    torch::Tensor X,
    torch::Tensor y,
    torch::Tensor X_test,
    float prior_precision,
    float noise_precision
);
"""

bayesian_extension = load_inline(
    name='bayesian_linear_v5',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['bayesian_linear_forward'],
    extra_cuda_cflags=['-O3', '-use_fast_math'],
    verbose=True
)


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.extension = bayesian_extension

    def forward(self, X: torch.Tensor, y: torch.Tensor, X_test: torch.Tensor) -> tuple:
        pred_mean, pred_var = self.extension.bayesian_linear_forward(
            X, y, X_test, self.prior_precision, self.noise_precision
        )
        return pred_mean, pred_var