import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Stein Variational Gradient Descent (SVGD) update step.
    Maintains a set of particles approximating a posterior and updates
    them via a kernel-smoothed gradient that balances fitting the target
    with repulsion between particles.  The RBF kernel bandwidth is set
    via the median heuristic.
    """

    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        # Target: N(0, I) for benchmarking
        self.register_buffer("target_mean", torch.zeros(dim))

    def _score_fn(self, x: torch.Tensor) -> torch.Tensor:
        """Gradient of log p(x) for N(0, I)."""
        return -x

    def forward(self, particles: torch.Tensor) -> torch.Tensor:
        """
        Args:
            particles: (N, dim) – current particle positions
        Returns:
            phi: (N, dim) – SVGD update direction
        """
        N = particles.shape[0]
        score = self._score_fn(particles)  # (N, dim)

        # Pairwise squared distances
        diff = particles.unsqueeze(0) - particles.unsqueeze(1)  # (N, N, dim)
        dist_sq = (diff ** 2).sum(-1)  # (N, N)

        # Median heuristic for bandwidth
        median_sq = torch.median(dist_sq.view(-1)).clamp(min=1e-5)
        h = median_sq / max(torch.log(torch.tensor(float(N))).item(), 1.0)

        K = torch.exp(-dist_sq / (2.0 * h))  # (N, N)
        grad_K = -diff / h * K.unsqueeze(-1)  # (N, N, dim)

        # SVGD update: phi(x_i) = (1/N) sum_j [ K(x_j, x_i) * score(x_j) + grad_K(x_j, x_i) ]
        phi = (K @ score + grad_K.sum(dim=0)) / N
        return phi


dim = 32
num_particles = 1024


def get_inputs():
    p = popcorn_pri
    return [torch.randn(num_particles, 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>
torch::Tensor svgd_forward(torch::Tensor particles);
"""

cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cmath>
#define HIST_BINS 32768
#define HIST_RANGE 256.0f

__global__ void dist_hist_kernel_D32(const float* __restrict__ x, float* __restrict__ dist, unsigned int* __restrict__ hist, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int M = N * N;
    if (idx >= M) return;
    int i = idx / N;
    int j = idx - i * N;
    const float* xi = x + i * 32;
    const float* xj = x + j * 32;
    float s = 0.0f;
    #pragma unroll
    for (int d = 0; d < 32; ++d) {
        float a = xj[d] - xi[d];
        s += a * a;
    }
    dist[idx] = s;
    int bin = (int)(s * ((float)HIST_BINS / HIST_RANGE));
    if (bin >= HIST_BINS) bin = HIST_BINS - 1;
    atomicAdd(hist + bin, 1u);
}

__global__ void dist_hist_kernel_gen(const float* __restrict__ x, float* __restrict__ dist, unsigned int* __restrict__ hist, int N, int D) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int M = N * N;
    if (idx >= M) return;
    int i = idx / N;
    int j = idx - i * N;
    const float* xi = x + i * D;
    const float* xj = x + j * D;
    float s = 0.0f;
    for (int d = 0; d < D; ++d) { float a = xj[d] - xi[d]; s += a*a; }
    dist[idx] = s;
    int bin = (int)(s * ((float)HIST_BINS / HIST_RANGE));
    if (bin >= HIST_BINS) bin = HIST_BINS - 1;
    if (bin < 0) bin = 0;
    atomicAdd(hist + bin, 1u);
}

__global__ void find_med_kernel(const unsigned int* __restrict__ hist, float* __restrict__ med, int M) {
    unsigned int target = (unsigned int)((M - 1) / 2);
    unsigned int cum = 0;
    for (int b = 0; b < HIST_BINS; ++b) {
        unsigned int c = hist[b];
        if (cum + c > target) { *med = ((float)b + 0.5f) * (HIST_RANGE / (float)HIST_BINS); return; }
        cum += c;
    }
    *med = HIST_RANGE;
}

__global__ void phi_kernel_D32(const float* __restrict__ x, const float* __restrict__ dist, const float* __restrict__ medptr, float* __restrict__ out, int N) {
    int i = blockIdx.x;
    int tid = threadIdx.x;
    float med = *medptr;
    if (med < 1.0e-5f) med = 1.0e-5f;
    float inv_h = logf((float)N) / med;
    if (N <= 2) inv_h = 1.0f / med;
    float inv_2h = 0.5f * inv_h;
    float inv_N = 1.0f / (float)N;
    float acc[32];
    #pragma unroll
    for (int d=0; d<32; ++d) acc[d] = 0.0f;
    const float* xi = x + i * 32;
    for (int j = tid; j < N; j += blockDim.x) {
        float K = expf(-dist[i * N + j] * inv_2h);
        const float* xj = x + j * 32;
        #pragma unroll
        for (int d=0; d<32; ++d) {
            float xjd = xj[d];
            acc[d] += K * (-xjd + (xjd - xi[d]) * inv_h);
        }
    }
    extern __shared__ float sh[];
    #pragma unroll
    for (int d=0; d<32; ++d) sh[d * blockDim.x + tid] = acc[d];
    __syncthreads();
    if (tid < 32) {
        float sum = 0.0f;
        for (int t=0; t<blockDim.x; ++t) sum += sh[tid * blockDim.x + t];
        out[i * 32 + tid] = sum * inv_N;
    }
}

__global__ void phi_kernel_generic(const float* x,const float* dist,const float* medptr,float* out,int N,int D) {
    int i=blockIdx.x, d=blockIdx.y, tid=threadIdx.x;
    float med=*medptr; if(med<1e-5f) med=1e-5f; float denom=logf((float)N); if(denom<1.0f) denom=1.0f;
    float inv_h=denom/med, inv_2h=0.5f*inv_h, xi=x[i*D+d], acc=0.0f;
    for(int j=tid;j<N;j+=blockDim.x){float xj=x[j*D+d]; float K=expf(-dist[i*N+j]*inv_2h); acc += K*(-xj+(xj-xi)*inv_h);} 
    extern __shared__ float sh[]; sh[tid]=acc; __syncthreads();
    for(int stride=blockDim.x/2;stride>0;stride>>=1){if(tid<stride) sh[tid]+=sh[tid+stride]; __syncthreads();}
    if(tid==0) out[i*D+d]=sh[0]/(float)N;
}

torch::Tensor svgd_forward(torch::Tensor particles) {
    TORCH_CHECK(particles.is_cuda(), "particles must be CUDA");
    TORCH_CHECK(particles.scalar_type() == torch::kFloat32, "particles must be float32");
    auto x = particles.contiguous();
    int N = (int)x.size(0), D = (int)x.size(1), M = N * N;
    auto dist = torch::empty({(long long)M}, x.options());
    auto out = torch::empty_like(x);
    auto hist = torch::empty({HIST_BINS}, torch::TensorOptions().device(x.device()).dtype(torch::kUInt32));
    auto med = torch::empty({1}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    C10_CUDA_CHECK(cudaMemsetAsync(hist.data_ptr<unsigned int>(), 0, HIST_BINS * sizeof(unsigned int), stream));
    int threads = 256, blocks = (M + threads - 1) / threads;
    if (D == 32) dist_hist_kernel_D32<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), dist.data_ptr<float>(), hist.data_ptr<unsigned int>(), N);
    else dist_hist_kernel_gen<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), dist.data_ptr<float>(), hist.data_ptr<unsigned int>(), N, D);
    find_med_kernel<<<1, 1, 0, stream>>>(hist.data_ptr<unsigned int>(), med.data_ptr<float>(), M);
    if (D == 32) {
        int t2 = 256;
        phi_kernel_D32<<<N, t2, (size_t)t2 * 32 * sizeof(float), stream>>>(x.data_ptr<float>(), dist.data_ptr<float>(), med.data_ptr<float>(), out.data_ptr<float>(), N);
    } else {
        dim3 grid(N, D); int t2 = 256;
        phi_kernel_generic<<<grid, t2, t2*sizeof(float), stream>>>(x.data_ptr<float>(), dist.data_ptr<float>(), med.data_ptr<float>(), out.data_ptr<float>(), N, D);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_ext = load_inline(name="svgd_ext_final", cpp_sources=cpp_src, cuda_sources=cuda_src, functions=["svgd_forward"], extra_cuda_cflags=["-O3", "--use_fast_math"], extra_cflags=["-O3"], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        self.register_buffer("target_mean", torch.zeros(dim))
    def forward(self, particles: torch.Tensor) -> torch.Tensor:
        return _ext.svgd_forward(particles)