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 L, double eps);
std::vector<torch::Tensor> hmc_forward(torch::Tensor q, torch::Tensor p0, torch::Tensor u, int64_t L, double eps) { return hmc_forward_cuda(q, p0, u, L, eps); }
"""

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

__device__ __forceinline__ float row_delta(float qi, float pi, float aq, float ap, float bq, float bp, float &qf) {
    qf = aq * qi + ap * pi;
    float pf = bq * qi + bp * pi;
    return 0.5f * ((qi * qi + pi * pi) - (qf * qf + pf * pf));
}

__global__ __launch_bounds__(32, 16) void hmc64_kernel(const float* __restrict__ q0,
        const float* __restrict__ p0, const float* __restrict__ u,
        float* __restrict__ qout, bool* __restrict__ acc,
        float aq, float ap, float bq, float bp) {
    int b = blockIdx.x;
    int lane = threadIdx.x;
    int idx0 = (b << 6) + lane;
    int idx1 = idx0 + 32;

    float qi0 = q0[idx0], pi0 = p0[idx0], qf0;
    float qi1 = q0[idx1], pi1 = p0[idx1], qf1;
    float delta = row_delta(qi0, pi0, aq, ap, bq, bp, qf0) + row_delta(qi1, pi1, aq, ap, bq, bp, qf1);

    unsigned mask = 0xffffffffu;
    delta += __shfl_down_sync(mask, delta, 16);
    delta += __shfl_down_sync(mask, delta, 8);
    delta += __shfl_down_sync(mask, delta, 4);
    delta += __shfl_down_sync(mask, delta, 2);
    delta += __shfl_down_sync(mask, delta, 1);

    int accept = 0;
    if (lane == 0) {
        accept = (logf(fmaxf(u[b], 1.0e-10f)) < delta);
        acc[b] = (accept != 0);
    }
    accept = __shfl_sync(mask, accept, 0);
    qout[idx0] = accept ? qf0 : qi0;
    qout[idx1] = accept ? qf1 : qi1;
}

__global__ __launch_bounds__(32, 16) void hmc_warp_kernel(const float* __restrict__ q0,
        const float* __restrict__ p0, const float* __restrict__ u,
        float* __restrict__ qout, bool* __restrict__ acc,
        int D, float aq, float ap, float bq, float bp) {
    int b = blockIdx.x;
    int lane = threadIdx.x;
    int base = b * D;
    float delta = 0.0f;
    for (int d = lane; d < D; d += 32) {
        float qi = q0[base + d], pi = p0[base + d], qf;
        delta += row_delta(qi, pi, aq, ap, bq, bp, qf);
    }
    unsigned mask = 0xffffffffu;
    delta += __shfl_down_sync(mask, delta, 16);
    delta += __shfl_down_sync(mask, delta, 8);
    delta += __shfl_down_sync(mask, delta, 4);
    delta += __shfl_down_sync(mask, delta, 2);
    delta += __shfl_down_sync(mask, delta, 1);

    int accept = 0;
    if (lane == 0) {
        accept = (logf(fmaxf(u[b], 1.0e-10f)) < delta);
        acc[b] = (accept != 0);
    }
    accept = __shfl_sync(mask, accept, 0);
    for (int d = lane; d < D; d += 32) {
        float qi = q0[base + d], pi = p0[base + d];
        float qf = aq * qi + ap * pi;
        qout[base + d] = accept ? qf : qi;
    }
}

std::vector<torch::Tensor> hmc_forward_cuda(torch::Tensor q, torch::Tensor p0, torch::Tensor u, int64_t L64, double epsd) {
    int B = (int)q.size(0), D = (int)q.size(1);
    auto qout = torch::empty_like(q);
    auto acc = torch::empty({B}, torch::TensorOptions().device(q.device()).dtype(torch::kBool));

    float eps = (float)epsd;
    float aq = 1.0f, ap = 0.0f, bq = -0.5f * eps, bp = 1.0f;
    int L = (int)L64;
    for (int i = 0; i < L - 1; ++i) {
        aq = aq + eps * bq;
        ap = ap + eps * bp;
        bq = bq - eps * aq;
        bp = bp - eps * ap;
    }
    aq = aq + eps * bq;
    ap = ap + eps * bp;
    bq = bq - 0.5f * eps * aq;
    bp = bp - 0.5f * eps * ap;

    cudaStream_t st = at::cuda::getCurrentCUDAStream();
    if (D == 64) {
        hmc64_kernel<<<B, 32, 0, st>>>(q.data_ptr<float>(), p0.data_ptr<float>(), u.data_ptr<float>(),
                                       qout.data_ptr<float>(), acc.data_ptr<bool>(), aq, ap, bq, bp);
    } else {
        hmc_warp_kernel<<<B, 32, 0, st>>>(q.data_ptr<float>(), p0.data_ptr<float>(), u.data_ptr<float>(),
                                          qout.data_ptr<float>(), acc.data_ptr<bool>(), D, aq, ap, bq, bp);
    }
    return {qout, acc};
}
"""

_ext = load_inline(name="hmc_inline_ext_submission", 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)
        return tuple(_ext.hmc_forward(q, p, u, int(self.num_leapfrog_steps), float(self.step_size)))