import torch
import torch.nn as nn


class Model(nn.Module):
    """
    One full sweep of Gibbs sampling for a multivariate Gaussian with
    known precision matrix.  Each coordinate is sampled from its full
    conditional distribution while holding all other coordinates fixed.
    Benchmarks the sequential-scan pattern common in Bayesian models
    with tractable conditionals (e.g., topic models, Boltzmann machines).
    """

    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        # Precision matrix: diagonally dominant for well-conditioning
        L = torch.randn(dim, dim) * 0.1
        precision = L @ L.t() + torch.eye(dim) * 2.0
        self.register_buffer("precision", precision)
        self.register_buffer("mu", torch.zeros(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, dim) – current state
        Returns:
            x_new: (B, dim) – state after one full Gibbs sweep
        """
        x = x.clone()
        for i in range(self.dim):
            prec_ii = self.precision[i, i]
            residual = x @ self.precision[i] - x[:, i] * prec_ii
            cond_mean = self.mu[i] - residual / prec_ii
            cond_var = 1.0 / prec_ii
            x[:, i] = cond_mean + torch.sqrt(cond_var) * torch.randn(x.shape[0], device=x.device)
        return x


dim = 32
batch_size = 4096


def get_inputs():
    return [torch.randn(popcorn_pri.jitter_int(batch_size), dim)]


def get_init_inputs():
    return [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>

__global__ void gibbs_kernel(
    const float* __restrict__ x_in,
    float* __restrict__ x_out,
    const float* __restrict__ precision,
    const float* __restrict__ rands,
    const float* __restrict__ mu,
    int batch_size,
    int dim
) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;

    for (int i = 0; i < dim; ++i) {
        float prec_ii = precision[i * dim + i];
        
        float residual = 0.0f;
        for (int j = 0; j < dim; ++j) {
            float xj = x_out[b * dim + j];
            residual += xj * precision[i * dim + j];
        }
        residual -= x_out[b * dim + i] * prec_ii;
        
        float cond_mean = mu[i] - residual / prec_ii;
        float cond_var = 1.0f / prec_ii;
        float std = sqrtf(cond_var);
        
        float rand_val = rands[b * dim + i];
        x_out[b * dim + i] = cond_mean + std * rand_val;
    }
}

torch::Tensor gibbs_cuda(
    torch::Tensor x,
    torch::Tensor precision,
    torch::Tensor mu,
    torch::Tensor rands
) {
    auto batch_size = x.size(0);
    auto dim = x.size(1);
    auto x_out = x.clone();
    
    const int threads = 256;
    const int blocks = (batch_size + threads - 1) / threads;
    
    gibbs_kernel<<<blocks, threads>>>(
        x.data_ptr<float>(),
        x_out.data_ptr<float>(),
        precision.data_ptr<float>(),
        rands.data_ptr<float>(),
        mu.data_ptr<float>(),
        batch_size,
        dim
    );
    
    return x_out;
}
'''

cpp_source = '''
torch::Tensor gibbs_cuda(
    torch::Tensor x,
    torch::Tensor precision,
    torch::Tensor mu,
    torch::Tensor rands
);
'''

module = load_inline(
    name='gibbs_sampler',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['gibbs_cuda'],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        L = torch.randn(dim, dim) * 0.1
        precision = L @ L.t() + torch.eye(dim) * 2.0
        self.register_buffer("precision", precision)
        self.register_buffer("mu", torch.zeros(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size = x.size(0)
        rands = torch.zeros(batch_size, self.dim, device=x.device)
        for i in range(self.dim):
            rands[:, i] = torch.randn(batch_size, device=x.device)
        return module.gibbs_cuda(x, self.precision, self.mu, rands)