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

# Define the CUDA kernel with warp-cooperative memory access
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

constexpr int MAX_DIM = 128;
constexpr int WARP_SIZE = 32;

template <typename scalar_t>
__global__ void __launch_bounds__(256) hmc_leapfrog_kernel(
    scalar_t* __restrict__ q_out,
    bool* __restrict__ accepted_out,
    const scalar_t* __restrict__ q_in,
    const scalar_t* __restrict__ p_init,
    const scalar_t* __restrict__ rand_uniform,
    const scalar_t* __restrict__ mu,
    int batch_size,
    int dim,
    int num_leapfrog_steps,
    scalar_t step_size
) {
    // Each thread handles one batch element
    int batch_idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // Shared memory for mu (broadcast to all threads)
    __shared__ scalar_t mu_shared[MAX_DIM];
    
    // Load mu cooperatively
    for (int d = threadIdx.x; d < dim; d += blockDim.x) {
        mu_shared[d] = mu[d];
    }
    __syncthreads();
    
    if (batch_idx >= batch_size) return;
    
    // Use vectorized loads for coalesced access
    // Each thread loads its portion of q and p
    scalar_t q[MAX_DIM];
    scalar_t p[MAX_DIM];
    scalar_t q_current[MAX_DIM];
    scalar_t p_current[MAX_DIM];
    
    // Load q and p using coalesced access
    // Threads access contiguous memory: thread i accesses q_in[i], q_in[batch_size+i], etc.
    // But we need to transpose the access for better coalescing
    
    // For now, use simple loads with __ldg for read-only cache
    const scalar_t* q_ptr = &q_in[batch_idx * dim];
    const scalar_t* p_ptr = &p_init[batch_idx * dim];
    
    // Use loop unrolling and ILP to hide latency
    #pragma unroll 8
    for (int d = 0; d < dim; d++) {
        q[d] = q_ptr[d];
        q_current[d] = q[d];
    }
    
    #pragma unroll 8
    for (int d = 0; d < dim; d++) {
        p[d] = p_ptr[d];
        p_current[d] = p[d];
    }
    
    scalar_t eps = step_size;
    scalar_t eps_half = 0.5 * eps;
    
    // Half step for momentum
    #pragma unroll 8
    for (int d = 0; d < dim; d++) {
        p[d] = p[d] + eps_half * (mu_shared[d] - q[d]);
    }
    
    // Full leapfrog steps
    for (int step = 0; step < num_leapfrog_steps - 1; step++) {
        #pragma unroll 8
        for (int d = 0; d < dim; d++) {
            q[d] = q[d] + eps * p[d];
        }
        #pragma unroll 8
        for (int d = 0; d < dim; d++) {
            p[d] = p[d] + eps * (mu_shared[d] - q[d]);
        }
    }
    
    // Final position update
    #pragma unroll 8
    for (int d = 0; d < dim; d++) {
        q[d] = q[d] + eps * p[d];
    }
    
    // Half step for momentum
    #pragma unroll 8
    for (int d = 0; d < dim; d++) {
        p[d] = p[d] + eps_half * (mu_shared[d] - q[d]);
    }
    
    // Negate momentum
    #pragma unroll 8
    for (int d = 0; d < dim; d++) {
        p[d] = -p[d];
    }
    
    // Compute Hamiltonians
    scalar_t current_H = 0.0;
    scalar_t proposed_H = 0.0;
    
    #pragma unroll 8
    for (int d = 0; d < dim; d++) {
        scalar_t diff_curr = q_current[d] - mu_shared[d];
        scalar_t diff_prop = q[d] - mu_shared[d];
        current_H += diff_curr * diff_curr + p_current[d] * p_current[d];
        proposed_H += diff_prop * diff_prop + p[d] * p[d];
    }
    current_H *= 0.5;
    proposed_H *= 0.5;
    
    // Accept/reject
    scalar_t log_alpha = current_H - proposed_H;
    scalar_t log_u = log(rand_uniform[batch_idx]);
    
    bool accept = log_u < log_alpha;
    accepted_out[batch_idx] = accept;
    
    // Write output
    scalar_t* q_out_ptr = &q_out[batch_idx * dim];
    #pragma unroll 8
    for (int d = 0; d < dim; d++) {
        q_out_ptr[d] = accept ? q[d] : q_current[d];
    }
}

std::tuple<torch::Tensor, torch::Tensor> hmc_leapfrog_cuda(
    torch::Tensor q_in,
    torch::Tensor p_init,
    torch::Tensor rand_uniform,
    torch::Tensor mu,
    int num_leapfrog_steps,
    double step_size
) {
    int batch_size = q_in.size(0);
    int dim = q_in.size(1);
    
    auto q_out = torch::empty_like(q_in);
    auto accepted_out = torch::empty({batch_size}, torch::TensorOptions().dtype(torch::kBool).device(q_in.device()));
    
    const int threads = 256;
    const int blocks = (batch_size + threads - 1) / threads;
    
    AT_DISPATCH_FLOATING_TYPES(q_in.scalar_type(), "hmc_leapfrog_kernel", ([&] {
        hmc_leapfrog_kernel<scalar_t><<<blocks, threads>>>(
            q_out.data_ptr<scalar_t>(),
            accepted_out.data_ptr<bool>(),
            q_in.data_ptr<scalar_t>(),
            p_init.data_ptr<scalar_t>(),
            rand_uniform.data_ptr<scalar_t>(),
            mu.data_ptr<scalar_t>(),
            batch_size,
            dim,
            num_leapfrog_steps,
            static_cast<scalar_t>(step_size)
        );
    }));
    
    return std::make_tuple(q_out, accepted_out);
}
"""

cpp_source = """
std::tuple<torch::Tensor, torch::Tensor> hmc_leapfrog_cuda(
    torch::Tensor q_in,
    torch::Tensor p_init,
    torch::Tensor rand_uniform,
    torch::Tensor mu,
    int num_leapfrog_steps,
    double step_size);
"""

# Load the CUDA extension
hmc_module = load_inline(
    name='hmc_extension_v4',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['hmc_leapfrog_cuda'],
    verbose=True,
)


class ModelNew(nn.Module):
    """
    Hamiltonian Monte Carlo (HMC) proposal using leapfrog integration.
    CUDA-optimized implementation.
    """

    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:
        """
        Args:
            q: (B, dim) – current positions
        Returns:
            (q_new, accepted):
                q_new:    (B, dim)
                accepted: (B,)
        """
        # Generate random momentum and uniform samples in PyTorch
        p = torch.randn_like(q)
        rand_uniform = torch.rand(q.shape[0], device=q.device, dtype=q.dtype)
        
        # Call CUDA kernel
        q_new, accepted = hmc_module.hmc_leapfrog_cuda(
            q, p, rand_uniform, self.mu, self.num_leapfrog_steps, self.step_size
        )
        
        return q_new, accepted