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

_cpp_src = r"""
#include <torch/extension.h>

torch::Tensor gibbs_forward_cuda(torch::Tensor x, torch::Tensor precision, torch::Tensor mu, torch::Tensor noise);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

#ifndef __CUDA_ARCH__
#define __CUDA_ARCH__ 0
#endif

__global__ void gibbs32_kernel(const float* __restrict__ x,
                               const float* __restrict__ P,
                               const float* __restrict__ mu,
                               const float* __restrict__ z,
                               float* __restrict__ out,
                               int B) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;

    const float* xin = x + ((long long)b) * 32;
    float v0 = xin[0];
    float v1 = xin[1];
    float v2 = xin[2];
    float v3 = xin[3];
    float v4 = xin[4];
    float v5 = xin[5];
    float v6 = xin[6];
    float v7 = xin[7];
    float v8 = xin[8];
    float v9 = xin[9];
    float v10 = xin[10];
    float v11 = xin[11];
    float v12 = xin[12];
    float v13 = xin[13];
    float v14 = xin[14];
    float v15 = xin[15];
    float v16 = xin[16];
    float v17 = xin[17];
    float v18 = xin[18];
    float v19 = xin[19];
    float v20 = xin[20];
    float v21 = xin[21];
    float v22 = xin[22];
    float v23 = xin[23];
    float v24 = xin[24];
    float v25 = xin[25];
    float v26 = xin[26];
    float v27 = xin[27];
    float v28 = xin[28];
    float v29 = xin[29];
    float v30 = xin[30];
    float v31 = xin[31];

    #pragma unroll
    for (int i = 0; i < 32; ++i) {
        const float* p = P + i * 32;
        float sum = 0.0f;
        sum = fmaf(v0,  p[0],  sum);
        sum = fmaf(v1,  p[1],  sum);
        sum = fmaf(v2,  p[2],  sum);
        sum = fmaf(v3,  p[3],  sum);
        sum = fmaf(v4,  p[4],  sum);
        sum = fmaf(v5,  p[5],  sum);
        sum = fmaf(v6,  p[6],  sum);
        sum = fmaf(v7,  p[7],  sum);
        sum = fmaf(v8,  p[8],  sum);
        sum = fmaf(v9,  p[9],  sum);
        sum = fmaf(v10, p[10], sum);
        sum = fmaf(v11, p[11], sum);
        sum = fmaf(v12, p[12], sum);
        sum = fmaf(v13, p[13], sum);
        sum = fmaf(v14, p[14], sum);
        sum = fmaf(v15, p[15], sum);
        sum = fmaf(v16, p[16], sum);
        sum = fmaf(v17, p[17], sum);
        sum = fmaf(v18, p[18], sum);
        sum = fmaf(v19, p[19], sum);
        sum = fmaf(v20, p[20], sum);
        sum = fmaf(v21, p[21], sum);
        sum = fmaf(v22, p[22], sum);
        sum = fmaf(v23, p[23], sum);
        sum = fmaf(v24, p[24], sum);
        sum = fmaf(v25, p[25], sum);
        sum = fmaf(v26, p[26], sum);
        sum = fmaf(v27, p[27], sum);
        sum = fmaf(v28, p[28], sum);
        sum = fmaf(v29, p[29], sum);
        sum = fmaf(v30, p[30], sum);
        sum = fmaf(v31, p[31], sum);

        float cur;
        if (i == 0) cur = v0;
        else if (i == 1) cur = v1;
        else if (i == 2) cur = v2;
        else if (i == 3) cur = v3;
        else if (i == 4) cur = v4;
        else if (i == 5) cur = v5;
        else if (i == 6) cur = v6;
        else if (i == 7) cur = v7;
        else if (i == 8) cur = v8;
        else if (i == 9) cur = v9;
        else if (i == 10) cur = v10;
        else if (i == 11) cur = v11;
        else if (i == 12) cur = v12;
        else if (i == 13) cur = v13;
        else if (i == 14) cur = v14;
        else if (i == 15) cur = v15;
        else if (i == 16) cur = v16;
        else if (i == 17) cur = v17;
        else if (i == 18) cur = v18;
        else if (i == 19) cur = v19;
        else if (i == 20) cur = v20;
        else if (i == 21) cur = v21;
        else if (i == 22) cur = v22;
        else if (i == 23) cur = v23;
        else if (i == 24) cur = v24;
        else if (i == 25) cur = v25;
        else if (i == 26) cur = v26;
        else if (i == 27) cur = v27;
        else if (i == 28) cur = v28;
        else if (i == 29) cur = v29;
        else if (i == 30) cur = v30;
        else cur = v31;

        float pii = p[i];
        float residual = sum - cur * pii;
        float nv = mu[i] - residual / pii + rsqrtf(pii) * z[(long long)i * B + b];

        if (i == 0) v0 = nv;
        else if (i == 1) v1 = nv;
        else if (i == 2) v2 = nv;
        else if (i == 3) v3 = nv;
        else if (i == 4) v4 = nv;
        else if (i == 5) v5 = nv;
        else if (i == 6) v6 = nv;
        else if (i == 7) v7 = nv;
        else if (i == 8) v8 = nv;
        else if (i == 9) v9 = nv;
        else if (i == 10) v10 = nv;
        else if (i == 11) v11 = nv;
        else if (i == 12) v12 = nv;
        else if (i == 13) v13 = nv;
        else if (i == 14) v14 = nv;
        else if (i == 15) v15 = nv;
        else if (i == 16) v16 = nv;
        else if (i == 17) v17 = nv;
        else if (i == 18) v18 = nv;
        else if (i == 19) v19 = nv;
        else if (i == 20) v20 = nv;
        else if (i == 21) v21 = nv;
        else if (i == 22) v22 = nv;
        else if (i == 23) v23 = nv;
        else if (i == 24) v24 = nv;
        else if (i == 25) v25 = nv;
        else if (i == 26) v26 = nv;
        else if (i == 27) v27 = nv;
        else if (i == 28) v28 = nv;
        else if (i == 29) v29 = nv;
        else if (i == 30) v30 = nv;
        else v31 = nv;
    }

    float* y = out + ((long long)b) * 32;
    y[0] = v0; y[1] = v1; y[2] = v2; y[3] = v3;
    y[4] = v4; y[5] = v5; y[6] = v6; y[7] = v7;
    y[8] = v8; y[9] = v9; y[10] = v10; y[11] = v11;
    y[12] = v12; y[13] = v13; y[14] = v14; y[15] = v15;
    y[16] = v16; y[17] = v17; y[18] = v18; y[19] = v19;
    y[20] = v20; y[21] = v21; y[22] = v22; y[23] = v23;
    y[24] = v24; y[25] = v25; y[26] = v26; y[27] = v27;
    y[28] = v28; y[29] = v29; y[30] = v30; y[31] = v31;
}

__global__ void gibbs_generic_kernel(const float* __restrict__ x,
                                     const float* __restrict__ P,
                                     const float* __restrict__ mu,
                                     const float* __restrict__ z,
                                     float* __restrict__ out,
                                     int B, int D) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;
    float vals[64];
    for (int j = 0; j < D && j < 64; ++j) vals[j] = x[(long long)b * D + j];
    for (int i = 0; i < D && i < 64; ++i) {
        const float* p = P + (long long)i * D;
        float sum = 0.0f;
        for (int j = 0; j < D && j < 64; ++j) sum = fmaf(vals[j], p[j], sum);
        float pii = p[i];
        float residual = sum - vals[i] * pii;
        vals[i] = mu[i] - residual / pii + rsqrtf(pii) * z[(long long)i * B + b];
    }
    for (int j = 0; j < D && j < 64; ++j) out[(long long)b * D + j] = vals[j];
}

torch::Tensor gibbs_forward_cuda(torch::Tensor x, torch::Tensor precision, torch::Tensor mu, torch::Tensor noise) {
    const int B = (int)x.size(0);
    const int D = (int)x.size(1);
    auto out = torch::empty_like(x);
    const int threads = 256;
    const int blocks = (B + threads - 1) / threads;
    if (D == 32) {
        gibbs32_kernel<<<blocks, threads>>>(x.data_ptr<float>(), precision.data_ptr<float>(), mu.data_ptr<float>(), noise.data_ptr<float>(), out.data_ptr<float>(), B);
    } else {
        gibbs_generic_kernel<<<blocks, threads>>>(x.data_ptr<float>(), precision.data_ptr<float>(), mu.data_ptr<float>(), noise.data_ptr<float>(), out.data_ptr<float>(), B, D);
    }
    return out;
}
"""

_gibbs_ext = load_inline(
    name="gibbs_sweep_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["gibbs_forward_cuda"],
    with_cuda=True,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    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:
        # Match the reference RNG stream: reference draws one contiguous (B,) vector
        # for each coordinate, so a contiguous (dim, B) tensor has the same order.
        noise = torch.randn((self.dim, x.shape[0]), device=x.device, dtype=x.dtype)
        return _gibbs_ext.gibbs_forward_cuda(x, self.precision, self.mu, noise)