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* x, const float* precision, const float* mu, const float* randn, float* out, int batch_size, int dim) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;

    float local_x[32];
    for (int i = 0; i < dim; i++) {
        local_x[i] = x[b * dim + i];
    }

    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++) {
            residual += local_x[j] * precision[i * dim + j];
        }
        residual -= local_x[i] * prec_ii;
        float cond_mean = mu[i] - residual / prec_ii;
        float cond_var = 1.0f / prec_ii;
        float std = sqrtf(cond_var);
        local_x[i] = cond_mean + std * randn[b * dim + i];
    }

    for (int i = 0; i < dim; i++) {
        out[b * dim + i] = local_x[i];
    }
}

torch::Tensor gibbs_cuda(torch::Tensor x, torch::Tensor precision, torch::Tensor mu, torch::Tensor randn) {
    int batch_size = x.size(0);
    int dim = x.size(1);
    auto out = torch::empty_like(x);
    dim3 blocks((batch_size + 255) / 256);
    gibbs_kernel<<<blocks, 256>>>(x.data_ptr<float>(), precision.data_ptr<float>(), mu.data_ptr<float>(), randn.data_ptr<float>(), out.data_ptr<float>(), batch_size, dim);
    return out;
}
'''

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

module = load_inline(
    name='gibbs_module',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['gibbs_cuda'],
    verbose=True,
    with_cuda=True
)

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:
        randn = torch.randn(x.shape[0], self.dim, device=x.device)
        return module.gibbs_cuda(x, self.precision, self.mu, randn)