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>

__device__ __forceinline__ float compute_log_prob(const float* __restrict__ x, const float* __restrict__ mu, const float* __restrict__ sigma, int dim) {
    float sum = 0.0f;
    for(int i = 0; i < dim; i++) {
        float val = (x[i] - mu[i]) / sigma[i];
        sum += val * val;
    }
    return -0.5f * sum;
}

__global__ void metropolis_kernel(
    const float* __restrict__ x,
    const float* __restrict__ proposal,
    const float* __restrict__ log_u,
    float* __restrict__ x_new,
    bool* __restrict__ accepted,
    const float* __restrict__ mu,
    const float* __restrict__ sigma,
    int B,
    int dim
) {
    int b = blockIdx.x;
    if (b >= B) return;

    extern __shared__ float smem[];
    float* prop_local = smem;
    float* x_local = smem + dim;

    for(int i = threadIdx.x; i < dim; i += blockDim.x) {
        x_local[i] = x[b * dim + i];
        prop_local[i] = proposal[b * dim + i];
    }
    __syncthreads();

    float lp_prop = compute_log_prob(prop_local, mu, sigma, dim);
    float lp_x = compute_log_prob(x_local, mu, sigma, dim);

    __shared__ bool acc_shared;
    if (threadIdx.x == 0) {
        float log_alpha = lp_prop - lp_x;
        acc_shared = log_u[b] < log_alpha;
        accepted[b] = acc_shared;
    }
    __syncthreads();

    for(int i = threadIdx.x; i < dim; i += blockDim.x) {
        x_new[b * dim + i] = acc_shared ? prop_local[i] : x_local[i];
    }
}

std::vector<torch::Tensor> metropolis_forward(
    torch::Tensor x,
    torch::Tensor proposal,
    torch::Tensor log_u,
    torch::Tensor mu,
    torch::Tensor sigma
) {
    int B = x.size(0);
    int dim = x.size(1);
    auto x_new = torch::empty_like(x);
    auto acc = torch::empty({B}, x.options().dtype(torch::kBool));

    int threads = 64;
    size_t shared = 2 * dim * sizeof(float);
    metropolis_kernel<<<B, threads, shared>>>(
        x.data_ptr<float>(),
        proposal.data_ptr<float>(),
        log_u.data_ptr<float>(),
        x_new.data_ptr<float>(),
        acc.data_ptr<bool>(),
        mu.data_ptr<float>(),
        sigma.data_ptr<float>(),
        B, dim
    );
    return {x_new, acc};
}
'''

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

module = load_inline(
    name='metropolis',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['metropolis_forward'],
    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):
        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))
        outputs = module.metropolis_forward(x, proposal, log_u, self.mu, self.sigma)
        return outputs[0], outputs[1]