import torch
import torch.nn as nn
import math


class Model(nn.Module):
    """
    Batched Metropolis-Hastings MCMC step.  Given a current state and a
    log-probability function (here a Gaussian target for benchmarking),
    proposes a new state via a Gaussian random walk, computes the
    acceptance ratio, and stochastically accepts or rejects.
    """

    def __init__(self, dim, step_size=0.1):
        super().__init__()
        self.dim = dim
        self.step_size = step_size
        # Target: N(mu, sigma^2 I)
        self.register_buffer("mu", torch.zeros(dim))
        self.register_buffer("sigma", torch.ones(dim))

    def _log_prob(self, x: torch.Tensor) -> torch.Tensor:
        return -0.5 * (((x - self.mu) / self.sigma) ** 2).sum(dim=-1)

    def forward(self, x: torch.Tensor) -> tuple:
        """
        Args:
            x: (B, dim) – current chain states
        Returns:
            (x_new, accepted):
                x_new:    (B, dim) – states after accept/reject
                accepted: (B,)    – boolean mask of accepted proposals
        """
        proposal = x + self.step_size * torch.randn_like(x)
        log_alpha = self._log_prob(proposal) - self._log_prob(x)
        log_u = torch.log(torch.rand(x.shape[0], device=x.device).clamp(min=1e-10))
        accepted = log_u < log_alpha
        x_new = torch.where(accepted.unsqueeze(-1), proposal, x)
        return x_new, accepted


dim = 64
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 = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <curand_kernel.h>

__global__ void mh_kernel(
    const float* __restrict__ x,
    float* __restrict__ x_new,
    bool* __restrict__ accepted,
    const float step_size,
    const int B,
    const int D,
    unsigned long long seed
) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;

    curandStatePhilox4_32_10_t state;
    curand_init(seed, b, 0, &state);

    float log_prob_x = 0.0f;
    float log_prob_prop = 0.0f;

    for (int d = 0; d < D; d++) {
        float val = x[b * D + d];
        log_prob_x += -0.5f * val * val;
    }

    for (int d = 0; d < D; d++) {
        float noise = curand_normal(&state);
        float prop = x[b * D + d] + step_size * noise;
        x_new[b * D + d] = prop;  // temporarily store proposal
        log_prob_prop += -0.5f * prop * prop;
    }

    float log_alpha = log_prob_prop - log_prob_x;
    float u = curand_uniform(&state);
    float log_u = logf(fmaxf(u, 1e-10f));
    bool acc = log_u < log_alpha;

    accepted[b] = acc;

    if (!acc) {
        for (int d = 0; d < D; d++) {
            x_new[b * D + d] = x[b * D + d];
        }
    }
}

std::vector<torch::Tensor> mh_cuda(torch::Tensor x, float step_size) {
    const int B = x.size(0);
    const int D = x.size(1);
    auto x_new = torch::empty_like(x);
    auto accepted = torch::empty({B}, x.options().dtype(torch::kBool));

    const int threads = 256;
    const int blocks = (B + threads - 1) / threads;

    unsigned long long seed = 123456789ULL;

    mh_kernel<<<blocks, threads>>>(
        x.data_ptr<float>(),
        x_new.data_ptr<float>(),
        accepted.data_ptr<bool>(),
        step_size,
        B,
        D,
        seed
    );

    return {x_new, accepted};
}
'''

_cpp_source = r'''
std::vector<torch::Tensor> mh_cuda(torch::Tensor x, float step_size);
'''

_mh_module = load_inline(
    name='mh_cuda',
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=['mh_cuda'],
    verbose=False
)

class ModelNew(nn.Module):
    def __init__(self, dim, step_size=0.1):
        super().__init__()
        self.dim = dim
        self.step_size = step_size
        self.register_buffer("mu", torch.zeros(dim))
        self.register_buffer("sigma", torch.ones(dim))

    def forward(self, x: torch.Tensor) -> tuple:
        result = _mh_module.mh_cuda(x.contiguous(), self.step_size)
        return result[0], result[1]
```