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

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

std::vector<torch::Tensor> mh_forward(torch::Tensor x, torch::Tensor noise, torch::Tensor u, double step_size);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#include <math.h>

__inline__ __device__ float warp_sum(float v) {
    unsigned mask = 0xffffffffu;
    v += __shfl_down_sync(mask, v, 16);
    v += __shfl_down_sync(mask, v, 8);
    v += __shfl_down_sync(mask, v, 4);
    v += __shfl_down_sync(mask, v, 2);
    v += __shfl_down_sync(mask, v, 1);
    return v;
}

// One warp handles one Markov chain row.  This is specialized for the benchmark
// shape (dim=64), but keeps a generic loop for any moderate dim.
__global__ void mh_kernel_f32(const float* __restrict__ x,
                              const float* __restrict__ noise,
                              const float* __restrict__ u,
                              float* __restrict__ out,
                              bool* __restrict__ acc,
                              int B, int D, float step) {
    int tid = threadIdx.x;
    int lane = tid & 31;
    int warp_in_block = tid >> 5;
    int row = blockIdx.x * 4 + warp_in_block;
    if (row >= B) return;

    int base = row * D;
    float sx = 0.0f;
    float sp = 0.0f;

    // Match PyTorch's step*noise followed by add as closely as possible by
    // forcing round-to-nearest multiply and add rather than a fused FMA.
    for (int d = lane; d < D; d += 32) {
        float xv = x[base + d];
        float nv = noise[base + d];
        float sn = __fmul_rn(step, nv);
        float pv = __fadd_rn(xv, sn);
        sx = __fadd_rn(sx, __fmul_rn(xv, xv));
        sp = __fadd_rn(sp, __fmul_rn(pv, pv));
    }

    sx = warp_sum(sx);
    sp = warp_sum(sp);

    float log_alpha = -0.5f * (sp - sx);
    float uu = u[row];
    uu = fmaxf(uu, 1.0e-10f);
    bool a = (logf(uu) < log_alpha);
    a = __shfl_sync(0xffffffffu, a, 0);

    if (lane == 0) acc[row] = a;

    for (int d = lane; d < D; d += 32) {
        float xv = x[base + d];
        if (a) {
            float nv = noise[base + d];
            float sn = __fmul_rn(step, nv);
            out[base + d] = __fadd_rn(xv, sn);
        } else {
            out[base + d] = xv;
        }
    }
}

std::vector<torch::Tensor> mh_forward(torch::Tensor x, torch::Tensor noise, torch::Tensor u, double step_size) {
    const int64_t B64 = x.size(0);
    const int64_t D64 = x.size(1);
    auto out = torch::empty_like(x);
    auto acc = torch::empty({B64}, x.options().dtype(torch::kBool));

    const int B = (int)B64;
    const int D = (int)D64;
    const dim3 block(128);
    const dim3 grid((B + 3) / 4);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    mh_kernel_f32<<<grid, block, 0, stream>>>(
        x.data_ptr<float>(), noise.data_ptr<float>(), u.data_ptr<float>(),
        out.data_ptr<float>(), acc.data_ptr<bool>(), B, D, (float)step_size);

    return {out, acc};
}
"""

_ext = load_inline(
    name="mh_mcmc_step_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["mh_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    with_cuda=True,
    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:
        # Generate random tensors with the same PyTorch RNG calls/order as the
        # reference implementation, so seeded correctness tests see identical
        # random variates.  The CUDA extension fuses proposal formation,
        # log-probability reduction, accept test, and select/copy.
        noise = torch.randn_like(x)
        u = torch.rand((x.shape[0],), device=x.device, dtype=x.dtype)
        return tuple(_ext.mh_forward(x, noise, u, float(self.step_size)))