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 = '''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#include <algorithm>
#include <cmath>

__global__ void compute_dist_sq(const float* particles, float* dist_sq, int N, int dim) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdx.y * blockDim.y + threadIdx.y;
    if (i < N && j < N) {
        float sum = 0.0f;
        for (int d = 0; d < dim; d++) {
            float diff = particles[i * dim + d] - particles[j * dim + d];
            sum += diff * diff;
        }
        dist_sq[i * N + j] = sum;
    }
}

__global__ void compute_phi(const float* particles, const float* score, const float* dist_sq, float* phi, float h, int N, int dim) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N) {
        float k_sum[32] = {0};
        float grad_sum[32] = {0};
        for (int j = 0; j < N; j++) {
            float ds = dist_sq[j * N + i];
            float K = expf(-ds / (2.0f * h));
            for (int d = 0; d < dim; d++) {
                float diff = particles[j * dim + d] - particles[i * dim + d];
                k_sum[d] += K * score[j * dim + d];
                grad_sum[d] += (diff / h) * K;
            }
        }
        for (int d = 0; d < dim; d++) {
            phi[i * dim + d] = (k_sum[d] + grad_sum[d]) / N;
        }
    }
}

torch::Tensor svgd_update(torch::Tensor particles) {
    int N = particles.size(0);
    int dim = particles.size(1);
    auto options = torch::TensorOptions().device(particles.device()).dtype(particles.dtype());
    auto dist_sq = torch::empty({N, N}, options);
    auto score = -particles;
    
    dim3 block(16, 16);
    dim3 grid((N + 15) / 16, (N + 15) / 16);
    compute_dist_sq<<<grid, block>>>(particles.data_ptr<float>(), dist_sq.data_ptr<float>(), N, dim);
    
    auto dist_cpu = dist_sq.cpu();
    float median_val = torch::median(dist_cpu.view(-1)).item<float>();
    median_val = std::max(median_val, 1e-5f);
    float logN = logf(static_cast<float>(N));
    float h = median_val / std::max(logN, 1.0f);
    
    auto phi = torch::empty({N, dim}, options);
    compute_phi<<<(N+255)/256, 256>>>(particles.data_ptr<float>(), score.data_ptr<float>(), dist_sq.data_ptr<float>(), phi.data_ptr<float>(), h, N, dim);
    return phi;
}
'''

cpp_source = '''
torch::Tensor svgd_update(torch::Tensor particles);
'''

module = load_inline(
    name='svgd_cuda',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['svgd_update'],
    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):
        return module.svgd_update(particles)