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_cuda(torch::Tensor particles);

torch::Tensor svgd_forward(torch::Tensor particles) {
    TORCH_CHECK(particles.is_cuda(), "particles must be a CUDA tensor");
    TORCH_CHECK(particles.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(particles.dim() == 2, "particles must be 2D");
    return svgd_forward_cuda(particles);
}
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <thrust/device_ptr.h>
#include <thrust/sort.h>
#include <thrust/system/cuda/execution_policy.h>
#include <math.h>

__global__ void pairwise_dist_kernel(const float* __restrict__ x,
                                     float* __restrict__ dist,
                                     int N, int D) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * N;
    if (idx >= total) 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;
    #pragma unroll
    for (int d = 0; d < 32; ++d) {
        if (d < D) {
            float v = xj[d] - xi[d];
            s += v * v;
        }
    }
    dist[idx] = s;
}

__global__ void phi_kernel(const float* __restrict__ x,
                           const float* __restrict__ sorted_dist,
                           float* __restrict__ out,
                           int N, int D, int kth, float logN) {
    int i = blockIdx.x;
    int lane = threadIdx.x;  // one warp, lane corresponds to dimension for D=32
    unsigned mask = 0xffffffffu;

    float med = sorted_dist[kth];
    med = med < 1.0e-5f ? 1.0e-5f : med;
    float h = med / logN;
    float inv_h = 1.0f / h;
    float inv_2h = 0.5f * inv_h;

    float xi = (lane < D) ? x[i * D + lane] : 0.0f;
    float acc = 0.0f;

    for (int j = 0; j < N; ++j) {
        float xj = (lane < D) ? x[j * D + lane] : 0.0f;
        float diff = xj - xi;
        float ss = (lane < D) ? diff * diff : 0.0f;

        // Sum squared distance over the warp.  D is 32 for the benchmark.
        ss += __shfl_down_sync(mask, ss, 16);
        ss += __shfl_down_sync(mask, ss, 8);
        ss += __shfl_down_sync(mask, ss, 4);
        ss += __shfl_down_sync(mask, ss, 2);
        ss += __shfl_down_sync(mask, ss, 1);
        float dist = __shfl_sync(mask, ss, 0);

        float K = expf(-dist * inv_2h);
        if (lane < D) {
            acc += K * (-xj + diff * inv_h);
        }
    }

    if (lane < D) {
        out[i * D + lane] = acc / (float)N;
    }
}

torch::Tensor svgd_forward_cuda(torch::Tensor particles) {
    auto x = particles.contiguous();
    const int N = (int)x.size(0);
    const int D = (int)x.size(1);
    auto out = torch::empty_like(x);
    auto dist = torch::empty({(long long)N * (long long)N}, x.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    int total = N * N;
    int threads = 256;
    int blocks = (total + threads - 1) / threads;
    pairwise_dist_kernel<<<blocks, threads, 0, stream>>>(
        x.data_ptr<float>(), dist.data_ptr<float>(), N, D);

    thrust::device_ptr<float> begin(dist.data_ptr<float>());
    thrust::device_ptr<float> end = begin + total;
    thrust::sort(thrust::cuda::par.on(stream), begin, end);

    int kth = (total - 1) / 2;  // torch.median returns the lower median for even counts
    float logN = logf((float)N);
    if (logN < 1.0f) logN = 1.0f;

    phi_kernel<<<N, 32, 0, stream>>>(
        x.data_ptr<float>(), dist.data_ptr<float>(), out.data_ptr<float>(), N, D, kth, logN);
    return out;
}
"""

_svgd_ext = load_inline(
    name="svgd_custom_ext_v1",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["svgd_forward"],
    extra_cflags=["-O3"],
    extra_cuda_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 _svgd_ext.svgd_forward(particles)