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

std::vector<torch::Tensor> hmc_forward(torch::Tensor q, torch::Tensor p0, torch::Tensor u, int64_t steps, double eps) {
    return hmc_forward_cuda(q, p0, u, steps, eps);
}
'''

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

__inline__ __device__ float warp_sum(float v) {
    v += __shfl_down_sync(0xffffffff, v, 16);
    v += __shfl_down_sync(0xffffffff, v, 8);
    v += __shfl_down_sync(0xffffffff, v, 4);
    v += __shfl_down_sync(0xffffffff, v, 2);
    v += __shfl_down_sync(0xffffffff, v, 1);
    return v;
}

__global__ void hmc64_kernel(const float* __restrict__ q_in,
                             const float* __restrict__ p0_in,
                             const float* __restrict__ u,
                             float* __restrict__ q_out,
                             bool* __restrict__ acc_out,
                             float aq, float bp, float cp, float dp) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    int idx = (b << 6) + tid;

    float q0 = q_in[idx];
    float p0 = p0_in[idx];
    float q = fmaf(bp, p0, aq * q0);
    float p = fmaf(dp, p0, cp * q0);
    float sum = 0.5f * (q0*q0 + p0*p0 - q*q - p*p);

    sum = warp_sum(sum);
    __shared__ float ws[2];
    __shared__ int accepted;
    if ((tid & 31) == 0) ws[tid >> 5] = sum;
    __syncthreads();

    if (tid == 0) {
        float s = ws[0] + ws[1];
        int a = (u[b] < __expf(s));
        accepted = a;
        acc_out[b] = (bool)a;
    }
    __syncthreads();
    q_out[idx] = accepted ? q : q0;
}

__global__ void hmc_generic_kernel(const float* __restrict__ q_in,
                                   const float* __restrict__ p0_in,
                                   const float* __restrict__ u,
                                   float* __restrict__ q_out,
                                   bool* __restrict__ acc_out,
                                   int D, float aq, float bp, float cp, float dp) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    extern __shared__ float smem[];
    float sum = 0.f, q0 = 0.f, p0 = 0.f, q = 0.f, p = 0.f;
    if (tid < D) {
        int idx = b * D + tid;
        q0 = q_in[idx];
        p0 = p0_in[idx];
        q = fmaf(bp, p0, aq * q0);
        p = fmaf(dp, p0, cp * q0);
        sum = 0.5f * (q0*q0 + p0*p0 - q*q - p*p);
    }
    smem[tid] = sum;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }
    bool acc = false;
    if (tid == 0) {
        acc = (u[b] < __expf(smem[0]));
        acc_out[b] = acc;
    }
    __syncthreads();
    acc = acc_out[b];
    if (tid < D) {
        int idx = b * D + tid;
        q_out[idx] = acc ? q : q0;
    }
}

std::vector<torch::Tensor> hmc_forward_cuda(torch::Tensor q, torch::Tensor p0, torch::Tensor u, int64_t steps64, double epsd) {
    int B = (int)q.size(0);
    int D = (int)q.size(1);

    float e = (float)epsd;
    float aq = 1.f, bp = 0.f, cp = -0.5f * e, dp = 1.f;
    for (int s = 0; s < (int)steps64 - 1; ++s) {
        float naq = fmaf(e, cp, aq);
        float nbp = fmaf(e, dp, bp);
        float ncp = fmaf(-e, naq, cp);
        float ndp = fmaf(-e, nbp, dp);
        aq = naq; bp = nbp; cp = ncp; dp = ndp;
    }
    aq = fmaf(e, cp, aq);
    bp = fmaf(e, dp, bp);
    cp = fmaf(-0.5f * e, aq, cp);
    dp = fmaf(-0.5f * e, bp, dp);

    auto qo = torch::empty_like(q);
    auto acc = torch::empty({B}, q.options().dtype(torch::kBool));
    cudaStream_t st = at::cuda::getCurrentCUDAStream();

    if (D == 64) {
        hmc64_kernel<<<B, 64, 0, st>>>(q.data_ptr<float>(), p0.data_ptr<float>(), u.data_ptr<float>(),
                                       qo.data_ptr<float>(), acc.data_ptr<bool>(), aq, bp, cp, dp);
    } else {
        int threads = 1;
        while (threads < D) threads <<= 1;
        if (threads < 32) threads = 32;
        if (threads > 1024) threads = 1024;
        hmc_generic_kernel<<<B, threads, threads * sizeof(float), st>>>(q.data_ptr<float>(), p0.data_ptr<float>(), u.data_ptr<float>(),
                                                                        qo.data_ptr<float>(), acc.data_ptr<bool>(), D, aq, bp, cp, dp);
    }
    return {qo, acc};
}
'''

_ext = load_inline(
    name='hmc_ext_final',
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=['hmc_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_cflags=['-O3'],
    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, dtype=q.dtype)
        return tuple(_ext.hmc_forward(q, p, u, int(self.num_leapfrog_steps), float(self.step_size)))