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_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#include <cmath>
#define ROWS_PER_BLOCK 8

__inline__ __device__ float warp_reduce_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;
}

__device__ __forceinline__ int accept_from_sum(float sum, float uu) {
    // If log_alpha >= 0, reference always accepts because log(rand) <= 0.
    if (sum <= 0.0f) return 1;
    if (uu < 1.0e-10f) uu = 1.0e-10f;
    return logf(uu) < (-0.5f * sum);
}

__global__ void mh_kernel64(const float* __restrict__ x, const float* __restrict__ eps, const float* __restrict__ u,
                            float* __restrict__ out, bool* __restrict__ acc,
                            int B, float step) {
    int tid = threadIdx.x;
    int lane = tid & 31;
    int warp = tid >> 5;
    int b = blockIdx.x * ROWS_PER_BLOCK + warp;
    if (b >= B) return;

    int base = b << 6;
    int idx0 = base + lane;
    float x0 = x[idx0];
    float p0 = fmaf(step, eps[idx0], x0);
    float sum = p0 * p0 - x0 * x0;

    int idx1 = idx0 + 32;
    float x1 = x[idx1];
    float p1 = fmaf(step, eps[idx1], x1);
    sum += p1 * p1 - x1 * x1;

    sum = warp_reduce_sum(sum);
    int ai = 0;
    if (lane == 0) {
        ai = accept_from_sum(sum, u[b]);
        acc[b] = (bool)ai;
    }
    ai = __shfl_sync(0xffffffffu, ai, 0);
    out[idx0] = ai ? p0 : x0;
    out[idx1] = ai ? p1 : x1;
}

__global__ void mh_kernel_generic(const float* __restrict__ x, const float* __restrict__ eps, const float* __restrict__ u,
                                  float* __restrict__ out, bool* __restrict__ acc,
                                  int B, int D, float step) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    extern __shared__ float sdata[];
    float sum = 0.0f;

    for (int d = tid; d < D; d += blockDim.x) {
        int idx = b * D + d;
        float xv = x[idx];
        float p = fmaf(step, eps[idx], xv);
        sum += p * p - xv * xv;
    }
    sdata[tid] = sum;
    __syncthreads();

    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (tid < stride) sdata[tid] += sdata[tid + stride];
        __syncthreads();
    }

    bool a = false;
    if (tid == 0) {
        a = accept_from_sum(sdata[0], u[b]);
        acc[b] = a;
        sdata[0] = a ? 1.0f : 0.0f;
    }
    __syncthreads();
    a = (sdata[0] != 0.0f);

    for (int d = tid; d < D; d += blockDim.x) {
        int idx = b * D + d;
        float xv = x[idx];
        float p = fmaf(step, eps[idx], xv);
        out[idx] = a ? p : xv;
    }
}

std::vector<torch::Tensor> mh_forward(torch::Tensor x, double step_size) {
    const int B = x.size(0);
    const int D = x.size(1);

    auto eps = torch::randn_like(x);
    auto u = torch::rand({B}, x.options());
    auto out = torch::empty_like(x);
    auto acc = torch::empty({B}, torch::TensorOptions().device(x.device()).dtype(torch::kBool));

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (D == 64) {
        mh_kernel64<<<(B + ROWS_PER_BLOCK - 1) / ROWS_PER_BLOCK, ROWS_PER_BLOCK * 32, 0, stream>>>(
            x.data_ptr<float>(), eps.data_ptr<float>(), u.data_ptr<float>(),
            out.data_ptr<float>(), acc.data_ptr<bool>(), B, (float)step_size);
    } else {
        int threads = (D <= 64) ? 64 : ((D <= 128) ? 128 : 256);
        mh_kernel_generic<<<B, threads, threads * sizeof(float), stream>>>(
            x.data_ptr<float>(), eps.data_ptr<float>(), u.data_ptr<float>(),
            out.data_ptr<float>(), acc.data_ptr<bool>(), B, D, (float)step_size);
    }
    return {out, acc};
}
'''

cpp_src = r'''
#include <torch/extension.h>
#include <vector>
std::vector<torch::Tensor> mh_forward(torch::Tensor x, double step_size);
'''

_ext = load_inline(
    name='mh_ext_submit',
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=['mh_forward'],
    verbose=False,
    extra_cuda_cflags=['-O3'],
    extra_cflags=['-O3']
)

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