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 <math.h>

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

    float sum_x2 = 0.0f;
    float sum_p2 = 0.0f;
    for (int i = 0; i < D; ++i) {
        float xv = x[b * D + i];
        float pv = proposal[b * D + i];
        sum_x2 += xv * xv;
        sum_p2 += pv * pv;
    }

    float log_alpha = -0.5f * (sum_p2 - sum_x2);
    bool acc = log_u[b] < log_alpha;

    accepted[b] = acc;
    for (int i = 0; i < D; ++i) {
        x_new[b * D + i] = acc ? proposal[b * D + i] : x[b * D + i];
    }
}

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

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

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

    return {x_new, accepted};
}
'''

cpp_source = r'''
std::vector<torch::Tensor> mh_step_cuda(torch::Tensor x, torch::Tensor proposal, torch::Tensor log_u);
'''

module = load_inline(
    name='mh_cuda',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['mh_step_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:
        proposal = x + self.step_size * torch.randn_like(x)
        log_u = torch.log(torch.rand(x.shape[0], device=x.device).clamp(min=1e-10))
        x_new, accepted = module.mh_step_cuda(x, proposal, log_u)
        return x_new, accepted
```