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

# Compile custom CUDA kernels at module load time
cuda_source = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>

// Constants for our problem
constexpr int BLOCK_SIZE = 256;
constexpr int D = 16;  // Input dimension

// Fused kernel: X^T @ X with scaling and adding diagonal
// Each block handles one batch element, computes DxD matrix
__global__ void compute_lambda_post_kernel(
    const float* __restrict__ X,  // (B, N, D)
    float* __restrict__ Lambda,    // (B, D, D)
    int B, int N,
    float prior_prec, float noise_prec
) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    
    // Shared memory for X tile and accumulation
    __shared__ float X_tile[256][16];  // Store X[b, :, :] transposed
    __shared__ float Lambda_shared[16][16];
    
    // Initialize Lambda_shared to zero
    for (int i = tid; i < D * D; i += blockDim.x) {
        int row = i / D;
        int col = i % D;
        Lambda_shared[row][col] = 0.0f;
    }
    __syncthreads();
    
    // Load X[b, :, :] into shared memory (transposed for coalesced access)
    const float* X_b = X + b * N * D;
    
    // Each thread processes multiple rows
    for (int n = tid; n < N; n += blockDim.x) {
        for (int d = 0; d < D; d++) {
            X_tile[n][d] = X_b[n * D + d];
        }
    }
    __syncthreads();
    
    // Compute X^T @ X using all threads
    // Each thread computes part of the result
    for (int row = 0; row < D; row++) {
        for (int col = row; col < D; col++) {
            float sum = 0.0f;
            for (int n = tid; n < N; n += blockDim.x) {
                sum += X_tile[n][row] * X_tile[n][col];
            }
            // Block-wide reduction
            for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
                sum += __shfl_down_sync(0xffffffff, sum, offset);
            }
            if (tid == 0) {
                Lambda_shared[row][col] = sum * noise_prec;
                Lambda_shared[col][row] = sum * noise_prec;  // Symmetric
            }
            __syncthreads();
        }
    }
    
    // Add diagonal prior
    if (tid < D) {
        Lambda_shared[tid][tid] += prior_prec;
    }
    __syncthreads();
    
    // Write result to global memory
    float* Lambda_b = Lambda + b * D * D;
    for (int i = tid; i < D * D; i += blockDim.x) {
        int row = i / D;
        int col = i % D;
        Lambda_b[row * D + col] = Lambda_shared[row][col];
    }
}

// Fused kernel: compute X^T @ y
__global__ void compute_xt_y_kernel(
    const float* __restrict__ X,  // (B, N, D)
    const float* __restrict__ y,  // (B, N, 1)
    float* __restrict__ Xty,       // (B, D, 1)
    int B, int N, float noise_prec
) {
    int b = blockIdx.x;
    int d = threadIdx.x;
    
    if (d >= D) return;
    
    const float* X_b = X + b * N * D;
    const float* y_b = y + b * N;
    
    float sum = 0.0f;
    for (int n = 0; n < N; n++) {
        sum += X_b[n * D + d] * y_b[n];
    }
    
    Xty[b * D + d] = sum * noise_prec;
}

// Fused kernel: compute predictive variance
// pred_var = sum_d sum_k (X_test * Sigma_post * X_test^T) + 1/noise_precision
__global__ void compute_pred_var_kernel(
    const float* __restrict__ X_test,     // (B, M, D)
    const float* __restrict__ Sigma_post, // (B, D, D)
    float* __restrict__ pred_var,          // (B, M)
    int B, int M, float inv_noise_prec
) {
    int b = blockIdx.y;
    int m = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (m >= M) return;
    
    const float* X_test_m = X_test + b * M * D + m * D;
    const float* Sigma = Sigma_post + b * D * D;
    
    float var = 0.0f;
    for (int d1 = 0; d1 < D; d1++) {
        float sum = 0.0f;
        for (int d2 = 0; d2 < D; d2++) {
            sum += Sigma[d1 * D + d2] * X_test_m[d2];
        }
        var += X_test_m[d1] * sum;
    }
    
    pred_var[b * M + m] = var + inv_noise_prec;
}

// Fused kernel for X_test @ mu_post (predictive mean)
__global__ void compute_pred_mean_kernel(
    const float* __restrict__ X_test,  // (B, M, D)
    const float* __restrict__ mu_post, // (B, D, 1)
    float* __restrict__ pred_mean,      // (B, M, 1)
    int B, int M
) {
    int b = blockIdx.y;
    int m = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (m >= M) return;
    
    const float* X_test_m = X_test + b * M * D + m * D;
    const float* mu = mu_post + b * D;
    
    float sum = 0.0f;
    for (int d = 0; d < D; d++) {
        sum += X_test_m[d] * mu[d];
    }
    
    pred_mean[b * M + m] = sum;
}

// Host functions
torch::Tensor compute_lambda_post_cuda(torch::Tensor X, float prior_prec, float noise_prec) {
    int B = X.size(0);
    int N = X.size(1);
    // int D = X.size(2);  // Known to be 16
    
    auto Lambda = torch::empty({B, D, D}, X.options());
    
    compute_lambda_post_kernel<<<B, BLOCK_SIZE>>>(
        X.data_ptr<float>(),
        Lambda.data_ptr<float>(),
        B, N, prior_prec, noise_prec
    );
    
    return Lambda;
}

torch::Tensor compute_xt_y_cuda(torch::Tensor X, torch::Tensor y, float noise_prec) {
    int B = X.size(0);
    int N = X.size(1);
    
    auto Xty = torch::empty({B, D, 1}, X.options());
    
    compute_xt_y_kernel<<<B, 32>>>(
        X.data_ptr<float>(),
        y.data_ptr<float>(),
        Xty.data_ptr<float>(),
        B, N, noise_prec
    );
    
    return Xty;
}

torch::Tensor compute_pred_mean_cuda(torch::Tensor X_test, torch::Tensor mu_post) {
    int B = X_test.size(0);
    int M = X_test.size(1);
    
    auto pred_mean = torch::empty({B, M, 1}, X_test.options());
    
    int threads = 64;
    int blocks_m = (M + threads - 1) / threads;
    dim3 grid(blocks_m, B);
    
    compute_pred_mean_kernel<<<grid, threads>>>(
        X_test.data_ptr<float>(),
        mu_post.data_ptr<float>(),
        pred_mean.data_ptr<float>(),
        B, M
    );
    
    return pred_mean;
}

torch::Tensor compute_pred_var_cuda(torch::Tensor X_test, torch::Tensor Sigma_post, float inv_noise_prec) {
    int B = X_test.size(0);
    int M = X_test.size(1);
    
    auto pred_var = torch::empty({B, M}, X_test.options());
    
    int threads = 64;
    int blocks_m = (M + threads - 1) / threads;
    dim3 grid(blocks_m, B);
    
    compute_pred_var_kernel<<<grid, threads>>>(
        X_test.data_ptr<float>(),
        Sigma_post.data_ptr<float>(),
        pred_var.data_ptr<float>(),
        B, M, inv_noise_prec
    );
    
    return pred_var;
}

// Fused kernel: X_test @ Sigma_post, stores result for variance computation
// Then compute variance directly
__global__ void compute_pred_var_fused_kernel(
    const float* __restrict__ X_test,     // (B, M, D)
    const float* __restrict__ Sigma_post, // (B, D, D)
    float* __restrict__ pred_var,          // (B, M)
    int B, int M, float inv_noise_prec
) {
    // Each block handles one (b, m) pair
    int b = blockIdx.y;
    int m = blockIdx.x;
    
    __shared__ float X_shared[16];  // One row of X_test
    __shared__ float S_shared[16][16];  // Sigma_post
    __shared__ float XS_shared[16];  // X_test @ Sigma_post result for row m
    
    // Collaborative load of X_test[b, m, :]
    for (int d = threadIdx.x; d < D; d += blockDim.x) {
        X_shared[d] = X_test[b * M * D + m * D + d];
    }
    
    // Collaborative load of Sigma_post[b, :, :]
    const float* Sigma = Sigma_post + b * D * D;
    for (int i = threadIdx.x; i < D * D; i += blockDim.x) {
        int row = i / D;
        int col = i % D;
        S_shared[row][col] = Sigma[i];
    }
    __syncthreads();
    
    // Compute X_test[m, :] @ Sigma_post
    // Each thread computes one element
    float local_xs[4] = {0.0f, 0.0f, 0.0f, 0.0f};
    for (int d = threadIdx.x; d < D; d += blockDim.x) {
        float sum = 0.0f;
        for (int k = 0; k < D; k++) {
            sum += X_shared[k] * S_shared[k][d];
        }
        if (threadIdx.x < D) {
            XS_shared[d] = sum;
        }
    }
    __syncthreads();
    
    // Compute sum over d of X_shared[d] * XS_shared[d]
    // This is (X @ Sigma @ X^T) for row m
    float dot = 0.0f;
    for (int d = threadIdx.x; d < D; d += blockDim.x) {
        dot += X_shared[d] * XS_shared[d];
    }
    
    // Reduce across threads
    for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
        dot += __shfl_down_sync(0xffffffff, dot, offset);
    }
    
    if (threadIdx.x == 0) {
        pred_var[b * M + m] = dot + inv_noise_prec;
    }
}

torch::Tensor compute_pred_var_fused_cuda(torch::Tensor X_test, torch::Tensor Sigma_post, float inv_noise_prec) {
    int B = X_test.size(0);
    int M = X_test.size(1);
    
    auto pred_var = torch::empty({B, M}, X_test.options());
    
    dim3 grid(M, B);  // One block per (b, m) pair
    int threads = 32;  // Enough threads for D=16
    
    compute_pred_var_fused_kernel<<<grid, threads>>>(
        X_test.data_ptr<float>(),
        Sigma_post.data_ptr<float>(),
        pred_var.data_ptr<float>(),
        B, M, inv_noise_prec
    );
    
    return pred_var;
}

// Combined kernel that does X^T @ X and X^T @ y together to reuse X in shared memory
__global__ void compute_xtx_and_xty_kernel(
    const float* __restrict__ X,  // (B, N, D)
    const float* __restrict__ y,  // (B, N, 1)
    float* __restrict__ Lambda,    // (B, D, D)
    float* __restrict__ Xty,       // (B, D, 1)
    int B, int N,
    float prior_prec, float noise_prec
) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    
    __shared__ float X_shared[256][16];  // X transposed in shared memory
    __shared__ float y_shared[256];       // y values
    __shared__ float Lambda_s[16][16];
    __shared__ float Xty_s[16];
    
    const float* X_b = X + b * N * D;
    const float* y_b = y + b * N;
    
    // Initialize outputs
    for (int i = tid; i < D * D; i += blockDim.x) {
        int row = i / D;
        int col = i % D;
        Lambda_s[row][col] = 0.0f;
    }
    if (tid < D) {
        Xty_s[tid] = 0.0f;
    }
    __syncthreads();
    
    // Load X and y into shared memory
    for (int n = tid; n < N; n += blockDim.x) {
        y_shared[n] = y_b[n];
        for (int d = 0; d < D; d++) {
            X_shared[n][d] = X_b[n * D + d];
        }
    }
    __syncthreads();
    
    // Compute X^T @ X and X^T @ y
    for (int row = 0; row < D; row++) {
        // Compute X^T @ y for this row
        float xty_sum = 0.0f;
        for (int n = tid; n < N; n += blockDim.x) {
            xty_sum += X_shared[n][row] * y_shared[n];
        }
        for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
            xty_sum += __shfl_down_sync(0xffffffff, xty_sum, offset);
        }
        if (tid == 0) {
            Xty_s[row] = xty_sum * noise_prec;
        }
        
        // Compute X^T @ X for columns >= row (upper triangle)
        for (int col = row; col < D; col++) {
            float xtx_sum = 0.0f;
            for (int n = tid; n < N; n += blockDim.x) {
                xtx_sum += X_shared[n][row] * X_shared[n][col];
            }
            for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
                xtx_sum += __shfl_down_sync(0xffffffff, xtx_sum, offset);
            }
            if (tid == 0) {
                float val = xtx_sum * noise_prec;
                Lambda_s[row][col] = val;
                Lambda_s[col][row] = val;
            }
        }
        __syncthreads();
    }
    
    // Add diagonal prior
    if (tid < D) {
        Lambda_s[tid][tid] += prior_prec;
    }
    __syncthreads();
    
    // Write Lambda to global memory
    float* Lambda_b = Lambda + b * D * D;
    for (int i = tid; i < D * D; i += blockDim.x) {
        int row = i / D;
        int col = i % D;
        Lambda_b[row * D + col] = Lambda_s[row][col];
    }
    
    // Write Xty to global memory
    float* Xty_b = Xty + b * D;
    if (tid < D) {
        Xty_b[tid] = Xty_s[tid];
    }
}

std::tuple<torch::Tensor, torch::Tensor> compute_xtx_and_xty_cuda(
    torch::Tensor X, torch::Tensor y, float prior_prec, float noise_prec
) {
    int B = X.size(0);
    int N = X.size(1);
    
    auto Lambda = torch::empty({B, D, D}, X.options());
    auto Xty = torch::empty({B, D, 1}, X.options());
    
    compute_xtx_and_xty_kernel<<<B, BLOCK_SIZE>>>(
        X.data_ptr<float>(),
        y.data_ptr<float>(),
        Lambda.data_ptr<float>(),
        Xty.data_ptr<float>(),
        B, N, prior_prec, noise_prec
    );
    
    return std::make_tuple(Lambda, Xty);
}
"""

cpp_source = r"""
std::tuple<torch::Tensor, torch::Tensor> compute_xtx_and_xty_cuda(torch::Tensor X, torch::Tensor y, float prior_prec, float noise_prec);
torch::Tensor compute_pred_mean_cuda(torch::Tensor X_test, torch::Tensor mu_post);
torch::Tensor compute_pred_var_fused_cuda(torch::Tensor X_test, torch::Tensor Sigma_post, float inv_noise_prec);
"""

# Load the extension
bayesian_ops = load_inline(
    name='bayesian_linear_cuda',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['compute_xtx_and_xty_cuda', 'compute_pred_mean_cuda', 'compute_pred_var_fused_cuda'],
    verbose=False,
    extra_cuda_cflags=['-O3', '--use_fast_math']
)


class ModelNew(nn.Module):
    """
    Bayesian linear regression with a conjugate Gaussian prior - optimized CUDA version.
    Uses custom kernels for the key operations:
    1. Fused X^T@X and X^T@y computation (shares X in shared memory)
    2. Fused predictive variance computation
    """

    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]
        
        # Ensure contiguous float32 tensors
        X = X.contiguous().float()
        y = y.contiguous().float()
        X_test = X_test.contiguous().float()
        
        # Fused kernel: compute X^T@X (Lambda_post scaled) and X^T@y (scaled)
        # Returns Lambda_post and Xty both scaled appropriately
        Lambda_post, rhs = bayesian_ops.compute_xtx_and_xty_cuda(
            X, y.squeeze(-1), self.prior_precision, self.noise_precision
        )
        rhs = rhs.unsqueeze(-1)  # (B, D, 1)
        
        # Cholesky decomposition and solve (use PyTorch optimized routines)
        L = torch.linalg.cholesky(Lambda_post)
        mu_post = torch.cholesky_solve(rhs, L)
        
        # Predictive mean using custom kernel
        pred_mean = bayesian_ops.compute_pred_mean_cuda(X_test, mu_post.squeeze(-1))
        pred_mean = pred_mean.unsqueeze(-1)  # (B, M, 1)
        
        # Predictive variance using fused kernel
        Sigma_post = torch.cholesky_inverse(L)
        pred_var = bayesian_ops.compute_pred_var_fused_cuda(
            X_test, Sigma_post, 1.0 / self.noise_precision
        )
        
        return pred_mean, pred_var