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, 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>

__global__ __launch_bounds__(32, 16)
void mh64_warp_kernel(const float* __restrict__ x,
                      const float* __restrict__ noise,
                      const float* __restrict__ u,
                      float* __restrict__ out,
                      bool* __restrict__ accepted,
                      float step) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int base = row * 64 + tid;

    float x0 = x[base];
    float p0 = x0 + step * noise[base];
    float x1 = x[base + 32];
    float p1 = x1 + step * noise[base + 32];
    float sum = (p0 * p0 - x0 * x0) + (p1 * p1 - x1 * x1);

    unsigned mask = 0xffffffffu;
    sum += __shfl_down_sync(mask, sum, 16);
    sum += __shfl_down_sync(mask, sum, 8);
    sum += __shfl_down_sync(mask, sum, 4);
    sum += __shfl_down_sync(mask, sum, 2);
    sum += __shfl_down_sync(mask, sum, 1);

    int a = 0;
    if (tid == 0) {
        float uu = u[row];
        if (uu < 1.0e-10f) uu = 1.0e-10f;
        a = (logf(uu) < (-0.5f * sum));
        accepted[row] = (bool)a;
    }
    a = __shfl_sync(mask, a, 0);
    out[base] = a ? p0 : x0;
    out[base + 32] = a ? p1 : x1;
}

__global__ void mh_generic_kernel(const float* __restrict__ x,
                                  const float* __restrict__ noise,
                                  const float* __restrict__ u,
                                  float* __restrict__ out,
                                  bool* __restrict__ accepted,
                                  int D, float step) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    float sum = 0.0f;
    for (int d = tid; d < D; d += blockDim.x) {
        int idx = row * D + d;
        float xi = x[idx];
        float p = xi + step * noise[idx];
        sum += p * p - xi * xi;
    }
    extern __shared__ float smem[];
    smem[tid] = sum;
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) smem[tid] += smem[tid + stride];
        __syncthreads();
    }
    __shared__ int acc_s;
    if (tid == 0) {
        float uu = u[row];
        if (uu < 1.0e-10f) uu = 1.0e-10f;
        int a = (logf(uu) < (-0.5f * smem[0]));
        acc_s = a;
        accepted[row] = (bool)a;
    }
    __syncthreads();
    for (int d = tid; d < D; d += blockDim.x) {
        int idx = row * D + d;
        float xi = x[idx];
        float p = xi + step * noise[idx];
        out[idx] = acc_s ? p : xi;
    }
}

std::vector<torch::Tensor> mh_forward(torch::Tensor x, double step_size) {
    TORCH_CHECK(x.is_cuda(), "x must be CUDA");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "only float32 supported");
    TORCH_CHECK(x.dim() == 2, "x must be 2D");
    x = x.contiguous();
    int B = (int)x.size(0);
    int D = (int)x.size(1);
    auto noise = torch::randn_like(x);
    auto u = torch::rand({B}, x.options());
    auto out = torch::empty_like(x);
    auto accepted = torch::empty({B}, torch::TensorOptions().device(x.device()).dtype(torch::kBool));
    float step = (float)step_size;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (D == 64) {
        mh64_warp_kernel<<<B, 32, 0, stream>>>(x.data_ptr<float>(), noise.data_ptr<float>(), u.data_ptr<float>(), out.data_ptr<float>(), accepted.data_ptr<bool>(), step);
    } else {
        int threads = 128;
        mh_generic_kernel<<<B, threads, threads * sizeof(float), stream>>>(x.data_ptr<float>(), noise.data_ptr<float>(), u.data_ptr<float>(), out.data_ptr<float>(), accepted.data_ptr<bool>(), D, step);
    }
    return {out, accepted};
}
'''

_ext = load_inline(name='mh_mcmc_ext_final', cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=['mh_forward'], extra_cuda_cflags=['-O3'], extra_cflags=['-O3'], 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:
        return tuple(_ext.mh_forward(x, float(self.step_size)))