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

svgd_cuda = load_inline(
    name='svgd_cuda',
    cpp_sources='''
#include <torch/extension.h>

torch::Tensor svgd_forward_cuda(torch::Tensor particles);
''',
    cuda_sources='''
#include <torch/extension.h>
#include <cfloat>
#include <thrust/sort.h>
#include <thrust/device_ptr.h>

// Block size for kernel launches
#define BLOCK_SIZE 256

// Kernel to compute pairwise squared distances
__global__ void compute_dist_sq_kernel(
    const float* __restrict__ particles,
    float* __restrict__ dist_sq,
    int N, int dim
) {
    int i = blockIdx.x;
    int j = threadIdx.x + blockIdx.y * blockDim.x;
    
    if (i < N && j < N) {
        float sum = 0.0f;
        for (int k = 0; k < dim; k++) {
            float diff = particles[i * dim + k] - particles[j * dim + k];
            sum += diff * diff;
        }
        dist_sq[i * N + j] = sum;
    }
}

// Kernel to compute diff, dist_sq, and accumulate for phi in one pass
__global__ void svgd_kernel(
    const float* __restrict__ particles,
    const float* __restrict__ score,
    float* __restrict__ phi,
    int N, int dim, float h
) {
    // Each block handles one output particle i
    int i = blockIdx.x;
    int tid = threadIdx.x;
    
    if (i >= N) return;
    
    // Shared memory for reduction
    extern __shared__ float sdata[];
    float* sphi = sdata;  // size: blockDim.x * dim
    
    // Initialize local accumulation
    float local_phi[32];  // max dim
    for (int k = 0; k < dim; k++) {
        local_phi[k] = 0.0f;
    }
    
    // Each thread handles multiple j values
    for (int j = tid; j < N; j += blockDim.x) {
        // Compute dist_sq and K
        float dist_sq = 0.0f;
        for (int k = 0; k < dim; k++) {
            float diff = particles[i * dim + k] - particles[j * dim + k];
            dist_sq += diff * diff;
        }
        
        float K = expf(-dist_sq / (2.0f * h));
        
        // Accumulate K * score[j] for each dimension
        for (int k = 0; k < dim; k++) {
            float diff_ik = particles[i * dim + k] - particles[j * dim + k];
            // phi[i] += K * score[j] + grad_K
            // grad_K = -diff / h * K
            local_phi[k] += K * score[j * dim + k] - diff_ik / h * K;
        }
    }
    
    // Store to shared memory
    for (int k = 0; k < dim; k++) {
        sphi[tid * dim + k] = local_phi[k];
    }
    __syncthreads();
    
    // Reduction in shared memory
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            for (int k = 0; k < dim; k++) {
                sphi[tid * dim + k] += sphi[(tid + stride) * dim + k];
            }
        }
        __syncthreads();
    }
    
    // Write result
    if (tid == 0) {
        for (int k = 0; k < dim; k++) {
            phi[i * dim + k] = sphi[k] / N;
        }
    }
}

torch::Tensor svgd_forward_cuda(torch::Tensor particles) {
    // Ensure contiguous and correct type
    particles = particles.contiguous().to(torch::kFloat32);
    
    int N = particles.size(0);
    int dim = particles.size(1);
    
    // Score function: -particles for N(0, I) target
    auto score = -particles;  // (N, dim)
    
    // Allocate dist_sq
    auto dist_sq = torch::empty({N, N}, particles.options());
    
    // Compute pairwise squared distances
    int blocks_x = N;
    int blocks_y = (N + BLOCK_SIZE - 1) / BLOCK_SIZE;
    dim3 grid(blocks_x, blocks_y);
    
    compute_dist_sq_kernel<<<grid, BLOCK_SIZE>>>(
        particles.data_ptr<float>(),
        dist_sq.data_ptr<float>(),
        N, dim
    );
    
    // Median heuristic using thrust::sort
    auto dist_flat = dist_sq.view(-1);
    thrust::device_ptr<float> thrust_ptr(dist_flat.data_ptr<float>());
    thrust::sort(thrust_ptr, thrust_ptr + N * N);
    
    // Get median value
    float median_sq = dist_flat[N * N / 2].item<float>();
    median_sq = std::max(median_sq, 1e-5f);
    
    float log_n = std::log(static_cast<float>(N));
    float h = median_sq / std::max(log_n, 1.0f);
    
    // Compute phi using SVGD kernel
    auto phi = torch::empty({N, dim}, particles.options());
    
    int threads_per_block = 128;  // must be power of 2 for reduction
    int shared_mem_size = threads_per_block * dim * sizeof(float);
    
    svgd_kernel<<<N, threads_per_block, shared_mem_size>>>(
        particles.data_ptr<float>(),
        score.data_ptr<float>(),
        phi.data_ptr<float>(),
        N, dim, h
    );
    
    return phi;
}
''',
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    functions=['svgd_forward_cuda']
)

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_cuda.svgd_forward_cuda(particles)