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

cuda_source = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void hmc_kernel(const float* __restrict__ q_in, const float* __restrict__ p_in, const float* __restrict__ u_in,
                           float* __restrict__ q_out, bool* __restrict__ acc_out,
                           int B, int D, float eps, int num_leapfrog) {
    extern __shared__ float smem[];
    int b = blockIdx.x;
    if (b >= B) return;
    float* q = smem;
    float* p = smem + D;
    float* cur_q = smem + 2 * D;
    for (int i = threadIdx.x; i < D; i += blockDim.x) {
        q[i] = q_in[b * D + i];
        cur_q[i] = q[i];
        p[i] = p_in[b * D + i];
    }
    __syncthreads();
    for (int i = threadIdx.x; i < D; i += blockDim.x) {
        p[i] += 0.5f * eps * (- q[i]);
    }
    __syncthreads();
    for (int s = 0; s < num_leapfrog - 1; s++) {
        for (int i = threadIdx.x; i < D; i += blockDim.x) q[i] += eps * p[i];
        __syncthreads();
        for (int i = threadIdx.x; i < D; i += blockDim.x) p[i] += eps * (-q[i]);
        __syncthreads();
    }
    for (int i = threadIdx.x; i < D; i += blockDim.x) q[i] += eps * p[i];
    __syncthreads();
    for (int i = threadIdx.x; i < D; i += blockDim.x) p[i] += 0.5f * eps * (-q[i]);
    __syncthreads();
    for (int i = threadIdx.x; i < D; i += blockDim.x) p[i] = -p[i];
    __syncthreads();
    if (threadIdx.x == 0) {
        float cH = 0.f, pH = 0.f;
        for(int i=0; i<D; i++) {
            cH += 0.5f * cur_q[i]*cur_q[i];
            pH += 0.5f * q[i]*q[i];
        }
        float la = cH - pH;
        float lu = logf(fmaxf(u_in[b], 1e-10f));
        bool acc = lu < la;
        acc_out[b] = acc;
        for(int i=0; i<D; i++) q_out[b*D+i] = acc ? q[i] : cur_q[i];
    }
}
'''

cpp_source = r'''
std::vector<torch::Tensor> launch_hmc(torch::Tensor q, torch::Tensor p, torch::Tensor u, int64_t steps, double eps) {
    int B = q.size(0);
    int D = q.size(1);
    auto qo = torch::empty_like(q);
    auto ao = torch::empty({B}, torch::TensorOptions().device(q.device()).dtype(torch::kBool));
    size_t sm = 3 * D * sizeof(float);
    hmc_kernel<<<B, 64, sm>>>(q.data_ptr<float>(), p.data_ptr<float>(), u.data_ptr<float>(), qo.data_ptr<float>(), ao.data_ptr<bool>(), B, D, (float)eps, (int)steps);
    return {qo, ao};
}
'''

_hmc_mod = load_inline(
    name="hmcmod",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["launch_hmc"],
    extra_cuda_cflags=["-O3", "--use_fast_math"]
)


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)
        outs = _hmc_mod.launch_hmc(q, p, u, self.num_leapfrog_steps, self.step_size)
        return outs[0], outs[1]