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__ log_u_in,
    float* __restrict__ q_out,
    bool* __restrict__ accepted_out,
    const int B, const int D, const float eps, const int num_leapfrog_steps
) {
    extern __shared__ float smem[];
    float* q = smem;
    float* p = smem + D;
    float* curr_q = smem + 2*D;
    float* curr_p = smem + 3*D;

    const int b = blockIdx.x;
    if (b >= B) return;
    const int tid = threadIdx.x;

    if (tid < D) {
        float qv = q_in[b * D + tid];
        q[tid] = qv;
        curr_q[tid] = qv;
        float pv = p_in[b * D + tid];
        p[tid] = pv;
        curr_p[tid] = pv;
    }
    __syncthreads();

    if (tid < D) p[tid] -= 0.5f * eps * q[tid];
    __syncthreads();

    for (int s = 0; s < num_leapfrog_steps - 1; ++s) {
        if (tid < D) q[tid] += eps * p[tid];
        __syncthreads();
        if (tid < D) p[tid] -= eps * q[tid];
        __syncthreads();
    }

    if (tid < D) q[tid] += eps * p[tid];
    __syncthreads();
    if (tid < D) p[tid] -= 0.5f * eps * q[tid];
    __syncthreads();

    if (tid < D) p[tid] = -p[tid];
    __syncthreads();

    if (tid == 0) {
        float sum_cq = 0.0f, sum_cp = 0.0f, sum_q = 0.0f, sum_p = 0.0f;
        for (int i = 0; i < D; ++i) {
            sum_cq += curr_q[i] * curr_q[i];
            sum_cp += curr_p[i] * curr_p[i];
            sum_q += q[i] * q[i];
            sum_p += p[i] * p[i];
        }
        float curr_H = 0.5f * sum_cq + 0.5f * sum_cp;
        float prop_H = 0.5f * sum_q + 0.5f * sum_p;
        float log_alpha = curr_H - prop_H;
        bool acc = log_u_in[b] < log_alpha;
        accepted_out[b] = acc;
        for (int i = 0; i < D; ++i) {
            q_out[b * D + i] = acc ? q[i] : curr_q[i];
        }
    }
}

std::tuple<torch::Tensor, torch::Tensor> launch_hmc(
    torch::Tensor q, torch::Tensor p, torch::Tensor log_u,
    int num_leapfrog_steps, float step_size, int dim
) {
    int B = q.size(0);
    auto q_new = torch::empty_like(q);
    auto accepted = torch::empty({B}, q.options().dtype(torch::kBool));
    int threads = dim;
    size_t smem_size = 4 * dim * sizeof(float);
    hmc_kernel<<<B, threads, smem_size>>>(
        q.data_ptr<float>(), p.data_ptr<float>(), log_u.data_ptr<float>(),
        q_new.data_ptr<float>(), accepted.data_ptr<bool>(),
        B, dim, step_size, num_leapfrog_steps
    );
    return std::make_tuple(q_new, accepted);
}
'''

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

std::tuple<torch::Tensor, torch::Tensor> launch_hmc(torch::Tensor q, torch::Tensor p, torch::Tensor log_u, int num_leapfrog_steps, float step_size, int dim);
'''

module = load_inline(
    name='hmc_cuda',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['launch_hmc'],
    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):
        p = torch.randn_like(q)
        current_H = 0.5 * (q ** 2).sum(dim=-1) + 0.5 * (p ** 2).sum(dim=-1)
        # but to match exactly, do full in kernel
        eps = self.step_size
        log_u = torch.log(torch.rand(q.shape[0], device=q.device).clamp(min=1e-10))
        q_new, accepted = module.launch_hmc(q.contiguous(), p, log_u, self.num_leapfrog_steps, eps, self.dim)
        return q_new, accepted