import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Hamiltonian Monte Carlo (HMC) proposal using leapfrog integration.
    Simulates Hamiltonian dynamics on a batched set of particles to
    produce proposals with high acceptance probability.  The target
    distribution is a standard Gaussian for benchmarking.
    """

    def __init__(self, dim, num_leapfrog_steps=10, step_size=0.1):
        super().__init__()
        self.dim = dim
        self.num_leapfrog_steps = num_leapfrog_steps
        self.step_size = step_size
        self.register_buffer("mu", torch.zeros(dim))

    def _grad_log_prob(self, x: torch.Tensor) -> torch.Tensor:
        """Gradient of log p(x) for N(0, I)."""
        return -(x - self.mu)

    def _log_prob(self, x: torch.Tensor) -> torch.Tensor:
        return -0.5 * (x ** 2).sum(dim=-1)

    def forward(self, q: torch.Tensor) -> tuple:
        """
        Args:
            q: (B, dim) – current positions
        Returns:
            (q_new, accepted):
                q_new:    (B, dim)
                accepted: (B,)
        """
        eps = self.step_size
        p = torch.randn_like(q)
        current_p = p.clone()
        current_q = q.clone()

        # Half step for momentum
        p = p + 0.5 * eps * self._grad_log_prob(q)

        for _ in range(self.num_leapfrog_steps - 1):
            q = q + eps * p
            p = p + eps * self._grad_log_prob(q)

        q = q + eps * p
        p = p + 0.5 * eps * self._grad_log_prob(q)
        p = -p  # negate for reversibility

        current_H = -self._log_prob(current_q) + 0.5 * (current_p ** 2).sum(-1)
        proposed_H = -self._log_prob(q) + 0.5 * (p ** 2).sum(-1)

        log_alpha = current_H - proposed_H
        log_u = torch.log(torch.rand(q.shape[0], device=q.device).clamp(min=1e-10))
        accepted = log_u < log_alpha
        q_new = torch.where(accepted.unsqueeze(-1), q, current_q)
        return q_new, accepted


dim = 64
batch_size = 2048


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> hmc_forward(torch::Tensor q,
                                       torch::Tensor p,
                                       torch::Tensor u,
                                       int64_t steps,
                                       double eps);
"""

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

__global__ void hmc_kernel(const float* __restrict__ q0,
                           const float* __restrict__ p0,
                           const float* __restrict__ u,
                           float* __restrict__ qout,
                           bool* __restrict__ accepted,
                           int B, int D, int steps, float eps) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    if (b >= B) return;

    const int base = b * D;
    float local = 0.0f;

    for (int d = tid; d < D; d += blockDim.x) {
        int idx = base + d;
        float q = q0[idx];
        float p = p0[idx];
        float cq = q;
        float cp = p;

        // Half momentum step: grad log p(q) = -q for N(0,I)
        p = p - (0.5f * eps) * q;

        // Full leapfrog steps except the last position update.
        for (int s = 0; s < steps - 1; ++s) {
            q = q + eps * p;
            p = p - eps * q;
        }

        q = q + eps * p;
        p = p - (0.5f * eps) * q;
        p = -p;

        qout[idx] = q;  // kept if accepted, overwritten below if rejected
        local += 0.5f * (cq * cq + cp * cp - q * q - p * p);
    }

    __shared__ float smem[256];
    smem[tid] = local;
    __syncthreads();

    // blockDim is 128 in the wrapper, but keep this generic for powers of two <=256.
    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[b];
        uu = uu < 1.0e-10f ? 1.0e-10f : uu;
        int acc = (logf(uu) < smem[0]);
        acc_s = acc;
        accepted[b] = (bool)acc;
    }
    __syncthreads();

    if (!acc_s) {
        for (int d = tid; d < D; d += blockDim.x) {
            int idx = base + d;
            qout[idx] = q0[idx];
        }
    }
}

std::vector<torch::Tensor> hmc_forward(torch::Tensor q,
                                       torch::Tensor p,
                                       torch::Tensor u,
                                       int64_t steps,
                                       double eps) {
    const int B = (int)q.size(0);
    const int D = (int)q.size(1);
    auto qout = torch::empty_like(q);
    auto acc = torch::empty({B}, q.options().dtype(torch::kBool));

    const int threads = 128;
    dim3 grid(B);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    hmc_kernel<<<grid, threads, 0, stream>>>(
        q.data_ptr<float>(), p.data_ptr<float>(), u.data_ptr<float>(),
        qout.data_ptr<float>(), acc.data_ptr<bool>(), B, D, (int)steps, (float)eps);
    return {qout, acc};
}
"""

_hmc_ext = load_inline(
    name="hmc_leapfrog_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["hmc_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, dim, num_leapfrog_steps=10, step_size=0.1):
        super().__init__()
        self.dim = dim
        self.num_leapfrog_steps = num_leapfrog_steps
        self.step_size = step_size
        self.register_buffer("mu", torch.zeros(dim))

    def forward(self, q: torch.Tensor) -> tuple:
        p = torch.randn_like(q)
        u = torch.rand((q.shape[0],), device=q.device)
        q_new, accepted = _hmc_ext.hmc_forward(q, p, u, int(self.num_leapfrog_steps), float(self.step_size))
        return q_new, accepted