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

cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void grad_sum_kernel(const float* __restrict__ particles, const float* __restrict__ K, float* __restrict__ grad_sum, int N, int dim, float h) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;
    for (int d = 0; d < dim; ++d) {
        float acc = 0.0f;
        for (int j = 0; j < N; ++j) {
            float diff = particles[j * dim + d] - particles[i * dim + d];
            float g = (diff / h) * K[j * N + i];
            acc += g;
        }
        grad_sum[i * dim + d] = acc;
    }
}

torch::Tensor grad_sum(torch::Tensor particles, torch::Tensor K, float h) {
    int N = particles.size(0);
    int dim = particles.size(1);
    auto grad_sum_t = torch::zeros({N, dim}, particles.options());
    int block1d = 256;
    int grid1d = (N + block1d - 1) / block1d;
    grad_sum_kernel<<<grid1d, block1d>>>(particles.data_ptr<float>(), K.data_ptr<float>(), grad_sum_t.data_ptr<float>(), N, dim, h);
    return grad_sum_t;
}
"""

cpp_source = """
torch::Tensor grad_sum(torch::Tensor particles, torch::Tensor K, float h);
"""

module = load_inline(
    name="svgd_cuda",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["grad_sum"],
    verbose=True,
    extra_cuda_cflags=["-std=c++17"],
)

class ModelNew(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        self.register_buffer("target_mean", torch.zeros(dim))
        self.module = module

    def forward(self, particles):
        N = particles.shape[0]
        score = -particles
        diff = particles.unsqueeze(0) - particles.unsqueeze(1)
        dist_sq = (diff ** 2).sum(-1)
        median_sq = torch.median(dist_sq.view(-1)).clamp(min=1e-5)
        h = median_sq / max(torch.log(torch.tensor(float(N), device=particles.device)).item(), 1.0)
        K = torch.exp(-dist_sq / (2.0 * h))
        score_term = K @ score
        grad_term = self.module.grad_sum(particles, K, h)
        phi = (score_term + grad_term) / N
        return phi