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>

// Constants
constexpr int D = 16;
constexpr int THREADS_PER_BLOCK = 256;

// Fused kernel for computing XtX and adding prior
// Each block handles one batch
__global__ void compute_lambda_post_kernel(
    const float* __restrict__ X,
    float* __restrict__ Lambda_post,
    int N, int B, float prior_prec, float noise_prec
) {
    int b = blockIdx.x;
    int tx = threadIdx.x;
    
    // X is (B, N, D), Lambda_post is (B, D, D)
    const float* Xb = X + b * N * D;
    float* Lb = Lambda_post + b * D * D;
    
    // Shared memory for partial sums
    __shared__ float s_XtX[D][D + 1];  // +1 to avoid bank conflicts
    
    // Initialize XtX to zero
    if (tx < D * D) {
        s_XtX[tx / D][tx % D] = 0.0f;
    }
    __syncthreads();
    
    // Each thread accumulates part of XtX
    // Process rows of X in chunks
    for (int iter = 0; iter < (N + 255) / 256; iter++) {
        int base_row = iter * 256;
        int valid_rows = min(256, N - base_row);
        
        // Each thread computes one or more entries of XtX
        for (int idx = tx; idx < D * D; idx += blockDim.x) {
            int i = idx / D;
            int j = idx % D;
            float sum = 0.0f;
            for (int k = 0; k < valid_rows; k++) {
                sum += Xb[(base_row + k) * D + i] * Xb[(base_row + k) * D + j];
            }
            atomicAdd(&s_XtX[i][j], sum);
        }
        __syncthreads();
    }
    
    // Store result with prior
    for (int idx = tx; idx < D * D; idx += blockDim.x) {
        int i = idx / D;
        int j = idx % D;
        float val = noise_prec * s_XtX[i][j];
        if (i == j) {
            val += prior_prec;
        }
        Lb[i * D + j] = val;
    }
}

// Cholesky decomposition for small matrices (D=16)
__global__ void cholesky_kernel(
    const float* __restrict__ A,
    float* __restrict__ L_out,
    int B
) {
    int b = blockIdx.x;
    const float* Ab = A + b * D * D;
    float* Lb = L_out + b * D * D;
    
    __shared__ float s_A[D][D + 1];
    __shared__ float s_L[D][D + 1];
    
    // Load A into shared memory
    if (threadIdx.x < D * D) {
        int row = threadIdx.x / D;
        int col = threadIdx.x % D;
        s_A[row][col] = Ab[row * D + col];
        s_L[row][col] = 0.0f;
    }
    __syncthreads();
    
    // Cholesky decomposition (single thread for correctness)
    if (threadIdx.x == 0) {
        for (int i = 0; i < D; i++) {
            // L[i][i] = sqrt(A[i][i] - sum(L[i][k]^2))
            float diag = s_A[i][i];
            for (int k = 0; k < i; k++) {
                diag -= s_L[i][k] * s_L[i][k];
            }
            s_L[i][i] = sqrtf(fmaxf(diag, 1e-10f));
            
            // L[j][i] = (A[j][i] - sum(L[j][k]*L[i][k])) / L[i][i]
            for (int j = i + 1; j < D; j++) {
                float val = s_A[j][i];
                for (int k = 0; k < i; k++) {
                    val -= s_L[j][k] * s_L[i][k];
                }
                s_L[j][i] = val / s_L[i][i];
            }
        }
    }
    __syncthreads();
    
    // Store L (copy entire matrix including upper triangle zeros)
    if (threadIdx.x < D * D) {
        int row = threadIdx.x / D;
        int col = threadIdx.x % D;
        Lb[row * D + col] = s_L[row][col];
    }
}

// Fused kernel: compute XtY and solve with Cholesky
__global__ void solve_mu_post_kernel(
    const float* __restrict__ X,
    const float* __restrict__ y,
    const float* __restrict__ L,
    float* __restrict__ mu_post,
    int N, int B, float noise_prec
) {
    int b = blockIdx.x;
    int tx = threadIdx.x;
    
    const float* Xb = X + b * N * D;
    const float* yb = y + b * N;
    const float* Lb = L + b * D * D;
    float* mb = mu_post + b * D;
    
    __shared__ float s_XtY[D];
    __shared__ float s_L[D][D + 1];
    __shared__ float s_rhs[D];
    __shared__ float s_z[D];
    __shared__ float s_mu[D];
    
    // Initialize
    if (tx < D) s_XtY[tx] = 0.0f;
    __syncthreads();
    
    // Compute XtY = X^T * y by accumulation
    for (int iter = 0; iter < (N + 255) / 256; iter++) {
        int base_row = iter * 256;
        int valid_rows = min(256, N - base_row);
        
        if (tx < D) {
            float sum = 0.0f;
            for (int k = 0; k < valid_rows; k++) {
                sum += Xb[(base_row + k) * D + tx] * yb[base_row + k];
            }
            atomicAdd(&s_XtY[tx], sum);
        }
        __syncthreads();
    }
    
    // Load L and compute rhs
    if (tx < D * D) {
        int row = tx / D;
        int col = tx % D;
        s_L[row][col] = Lb[row * D + col];
    }
    if (tx < D) {
        s_rhs[tx] = noise_prec * s_XtY[tx];
    }
    __syncthreads();
    
    // Solve L * L^T * mu = rhs
    // Step 1: L * z = rhs (forward substitution)
    if (tx == 0) {
        for (int i = 0; i < D; i++) {
            float sum = s_rhs[i];
            for (int j = 0; j < i; j++) {
                sum -= s_L[i][j] * s_z[j];
            }
            s_z[i] = sum / s_L[i][i];
        }
        
        // Step 2: L^T * mu = z (backward substitution)
        for (int i = D - 1; i >= 0; i--) {
            float sum = s_z[i];
            for (int j = i + 1; j < D; j++) {
                sum -= s_L[j][i] * s_mu[j];
            }
            s_mu[i] = sum / s_L[i][i];
        }
    }
    __syncthreads();
    
    // Store mu_post
    if (tx < D) {
        mb[tx] = s_mu[tx];
    }
}

// Predictive mean: pred_mean = X_test @ mu_post
__global__ void predict_mean_kernel(
    const float* __restrict__ X_test,
    const float* __restrict__ mu_post,
    float* __restrict__ pred_mean,
    int M, int B
) {
    int b = blockIdx.x;
    int tx = threadIdx.x;
    
    const float* Xtb = X_test + b * M * D;
    const float* mb = mu_post + b * D;
    float* pb = pred_mean + b * M;
    
    __shared__ float s_mu[D];
    
    // Load mu_post into shared memory
    if (tx < D) {
        s_mu[tx] = mb[tx];
    }
    __syncthreads();
    
    // Each thread computes one or more predictions
    for (int i = tx; i < M; i += blockDim.x) {
        float sum = 0.0f;
        for (int j = 0; j < D; j++) {
            sum += Xtb[i * D + j] * s_mu[j];
        }
        pb[i] = sum;
    }
}

// Predictive variance using L^{-1}
__global__ void predict_var_kernel(
    const float* __restrict__ X_test,
    const float* __restrict__ L,
    float* __restrict__ pred_var,
    int M, int B, float noise_prec
) {
    int b = blockIdx.x;
    int tx = threadIdx.x;
    
    const float* Xtb = X_test + b * M * D;
    const float* Lb = L + b * D * D;
    float* vb = pred_var + b * M;
    
    __shared__ float s_L[D][D + 1];
    __shared__ float s_Linv[D][D + 1];
    
    // Load L
    if (tx < D * D) {
        int row = tx / D;
        int col = tx % D;
        s_L[row][col] = Lb[row * D + col];
        s_Linv[row][col] = 0.0f;
    }
    __syncthreads();
    
    // Compute L^{-1} by forward substitution: L * L_inv = I
    // For each column col of L_inv, solve L * (L_inv[:,col]) = e_col
    if (tx == 0) {
        for (int col = 0; col < D; col++) {
            for (int i = 0; i < D; i++) {
                // We need L_inv[i,col], but L_inv is lower triangular so only compute for i >= col
                if (i >= col) {
                    float rhs = (i == col) ? 1.0f : 0.0f;
                    float sum = rhs;
                    for (int j = 0; j < i; j++) {
                        sum -= s_L[i][j] * s_Linv[j][col];
                    }
                    s_Linv[i][col] = sum / s_L[i][i];
                }
                // For i < col, s_Linv[i][col] stays 0 (upper triangular part)
            }
        }
    }
    __syncthreads();
    
    // pred_var[i] = x_i @ Sigma @ x_i^T = ||L^{-1} @ x_i^T||^2
    // where x_i is the i-th row of X_test
    for (int i = tx; i < M; i += blockDim.x) {
        float result[D];
        
        // Compute L_inv @ x^T
        // result[row] = sum_{col=0}^{row} s_Linv[row][col] * x[col]
        for (int row = 0; row < D; row++) {
            float sum = 0.0f;
            for (int col = 0; col <= row; col++) {  // L_inv is lower triangular
                sum += s_Linv[row][col] * Xtb[i * D + col];
            }
            result[row] = sum;
        }
        
        // Compute squared norm
        float var = 0.0f;
        for (int j = 0; j < D; j++) {
            var += result[j] * result[j];
        }
        vb[i] = var + 1.0f / noise_prec;
    }
}

// Host function
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 D = X.size(2);
    int M = X_test.size(1);
    
    auto options = X.options();
    
    // Allocate intermediate tensors
    auto Lambda_post = torch::empty({B, D, D}, options);
    auto L = torch::empty({B, D, D}, options);
    auto mu_post = torch::empty({B, D}, options);
    auto pred_mean = torch::empty({B, M}, options);
    auto pred_var = torch::empty({B, M}, options);
    
    // Launch kernels
    int blocks = B;
    
    // Compute Lambda_post = prior_precision * I + noise_precision * X^T @ X
    compute_lambda_post_kernel<<<blocks, THREADS_PER_BLOCK>>>(
        X.data_ptr<float>(),
        Lambda_post.data_ptr<float>(),
        N, B, prior_precision, noise_precision
    );
    
    // Cholesky decomposition
    cholesky_kernel<<<blocks, THREADS_PER_BLOCK>>>(
        Lambda_post.data_ptr<float>(),
        L.data_ptr<float>(),
        B
    );
    
    // Solve for mu_post
    solve_mu_post_kernel<<<blocks, THREADS_PER_BLOCK>>>(
        X.data_ptr<float>(),
        y.data_ptr<float>(),
        L.data_ptr<float>(),
        mu_post.data_ptr<float>(),
        N, B, noise_precision
    );
    
    // Predictive mean
    predict_mean_kernel<<<blocks, THREADS_PER_BLOCK>>>(
        X_test.data_ptr<float>(),
        mu_post.data_ptr<float>(),
        pred_mean.data_ptr<float>(),
        M, B
    );
    
    // Predictive variance
    predict_var_kernel<<<blocks, THREADS_PER_BLOCK>>>(
        X_test.data_ptr<float>(),
        L.data_ptr<float>(),
        pred_var.data_ptr<float>(),
        M, B, noise_precision
    );
    
    // Reshape outputs to match expected shapes
    return std::make_tuple(
        pred_mean.unsqueeze(-1),  // (B, M, 1)
        pred_var  // (B, M)
    );
}
"""

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_cuda_v2',
    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
    
    def forward(self, X: torch.Tensor, y: torch.Tensor, X_test: torch.Tensor):
        pred_mean, pred_var = bayesian_extension.bayesian_linear_forward(
            X.contiguous(),
            y.contiguous(),
            X_test.contiguous(),
            self.prior_precision,
            self.noise_precision
        )
        return pred_mean, pred_var