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
import math

cuda_source = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <thrust/device_vector.h>
#include <thrust/sort.h>
#include <c10/cuda/CUDAStream.h>

// Block sizes
#define BLOCK_SIZE 256
#define TILE_DIM 32

// Compute pairwise squared distances and store in row-major
__global__ void compute_dist_sq_kernel(
    const float* __restrict__ particles,
    float* __restrict__ dist_sq,
    int N, int dim
) {
    int i = blockIdx.y;
    int j = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (i >= N || j >= N) return;
    
    float sum = 0.0f;
    for (int k = 0; k < dim; k++) {
        float diff = particles[j * dim + k] - particles[i * dim + k];
        sum += diff * diff;
    }
    dist_sq[i * N + j] = sum;
}

// Compute score = -particles
__global__ void compute_score_kernel(
    const float* __restrict__ particles,
    float* __restrict__ score,
    int N, int dim
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * dim;
    
    if (idx < total) {
        score[idx] = -particles[idx];
    }
}

// Compute kernel matrix K = exp(-dist_sq / (2*h))
__global__ void compute_kernel_kernel(
    const float* __restrict__ dist_sq,
    float* __restrict__ K,
    float inv_2h,
    int N
) {
    int i = blockIdx.y;
    int j = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (i >= N || j >= N) return;
    
    float d = dist_sq[i * N + j];
    K[i * N + j] = expf(-d * inv_2h);
}

// Compute phi = (K @ score + grad_K.sum(dim=0)) / N
// grad_K[i,j,k] = -(particles[j,k] - particles[i,k]) / h * K[i,j]
// phi1[j,k] = sum_i K[i,j] * score[j,k] = score[j,k] * sum_i K[i,j]
// phi2[j,k] = sum_i grad_K[i,j,k] = sum_i -(particles[j,k] - particles[i,k]) / h * K[i,j]
//          = (1/h) * sum_i (particles[i,k] - particles[j,k]) * K[i,j]
//          = (1/h) * (sum_i particles[i,k] * K[i,j] - particles[j,k] * sum_i K[i,j])

__global__ void compute_phi_kernel(
    const float* __restrict__ particles,
    const float* __restrict__ K,
    const float* __restrict__ score,
    float* __restrict__ phi,
    float inv_h,
    int N, int dim
) {
    int j = blockIdx.x;
    int k = blockIdx.y;
    
    if (j >= N || k >= dim) return;
    
    float sum_K_col = 0.0f;
    float sum_pK = 0.0f;
    float K_score_sum = 0.0f;
    
    float score_jk = score[j * dim + k];
    float p_jk = particles[j * dim + k];
    
    // Iterate over i (rows of K)
    // K[i,j] is at K[i*N + j]
    for (int i = 0; i < N; i++) {
        float Kij = K[i * N + j];
        sum_K_col += Kij;
        sum_pK += particles[i * dim + k] * Kij;
        K_score_sum += Kij * score[i * dim + k];
    }
    
    // phi1[j,k] = sum_i K[i,j] * score[i,k] = K_score_sum
    // phi2[j,k] = (1/h) * (sum_pK - p_jk * sum_K_col)
    float phi1 = K_score_sum;
    float phi2 = inv_h * (sum_pK - p_jk * sum_K_col);
    
    phi[j * dim + k] = (phi1 + phi2) / (float)N;
}

// Optimized version using shared memory for reduction
__global__ void compute_phi_optimized_kernel(
    const float* __restrict__ particles,
    const float* __restrict__ K,
    const float* __restrict__ score,
    float* __restrict__ phi,
    float inv_h,
    int N, int dim
) {
    extern __shared__ float shared_data[];
    
    int j = blockIdx.x;
    int k = blockIdx.y;
    
    if (j >= N || k >= dim) return;
    
    int tid = threadIdx.x;
    int blockSize = blockDim.x;
    
    float score_jk = score[j * dim + k];
    float p_jk = particles[j * dim + k];
    
    float local_K_score = 0.0f;
    float local_sum_K_col = 0.0f;
    float local_sum_pK = 0.0f;
    
    // Strided loop over i
    for (int i = tid; i < N; i += blockSize) {
        float Kij = K[i * N + j];
        float pik = particles[i * dim + k];
        float sik = score[i * dim + k];
        
        local_K_score += Kij * sik;
        local_sum_K_col += Kij;
        local_sum_pK += pik * Kij;
    }
    
    // Shared memory for reduction
    float* s_K_score = shared_data;
    float* s_sum_K = shared_data + blockSize;
    float* s_sum_pK = shared_data + 2 * blockSize;
    
    s_K_score[tid] = local_K_score;
    s_sum_K[tid] = local_sum_K_col;
    s_sum_pK[tid] = local_sum_pK;
    
    __syncthreads();
    
    // Block reduction
    for (int s = blockSize / 2; s > 0; s >>= 1) {
        if (tid < s) {
            s_K_score[tid] += s_K_score[tid + s];
            s_sum_K[tid] += s_sum_K[tid + s];
            s_sum_pK[tid] += s_sum_pK[tid + s];
        }
        __syncthreads();
    }
    
    if (tid == 0) {
        float phi1 = s_K_score[0];
        float phi2 = inv_h * (s_sum_pK[0] - p_jk * s_sum_K[0]);
        phi[j * dim + k] = (phi1 + phi2) / (float)N;
    }
}

torch::Tensor svgd_forward_cuda(torch::Tensor particles) {
    int N = particles.size(0);
    int dim = particles.size(1);
    
    auto options = particles.options();
    
    // Allocate intermediate tensors
    auto dist_sq = torch::empty({N, N}, options);
    auto K = torch::empty({N, N}, options);
    auto score = torch::empty({N, dim}, options);
    auto phi = torch::empty({N, dim}, options);
    
    // Compute score = -particles
    int total_threads = N * dim;
    int blocks = (total_threads + BLOCK_SIZE - 1) / BLOCK_SIZE;
    compute_score_kernel<<<blocks, BLOCK_SIZE>>>(
        particles.data_ptr<float>(),
        score.data_ptr<float>(),
        N, dim
    );
    
    // Compute pairwise squared distances
    dim3 grid((N + BLOCK_SIZE - 1) / BLOCK_SIZE, N);
    compute_dist_sq_kernel<<<grid, BLOCK_SIZE>>>(
        particles.data_ptr<float>(),
        dist_sq.data_ptr<float>(),
        N, dim
    );
    
    // Find median using Thrust sort
    auto dist_sq_flat = dist_sq.view(-1);
    thrust::device_ptr<float> d_ptr(dist_sq_flat.data_ptr<float>());
    thrust::sort(d_ptr, d_ptr + N * N);
    
    float median_sq;
    cudaMemcpy(&median_sq, dist_sq_flat.data_ptr<float>() + (N * N) / 2, sizeof(float), cudaMemcpyDeviceToHost);
    
    // Clamp median_sq to minimum
    if (median_sq < 1e-5f) median_sq = 1e-5f;
    
    // Compute h
    float log_N = logf((float)N);
    if (log_N < 1.0f) log_N = 1.0f;
    float h = median_sq / log_N;
    float inv_2h = 0.5f / h;
    float inv_h = 1.0f / h;
    
    // Compute kernel matrix
    compute_kernel_kernel<<<grid, BLOCK_SIZE>>>(
        dist_sq.data_ptr<float>(),
        K.data_ptr<float>(),
        inv_2h,
        N
    );
    
    // Compute phi using optimized kernel
    dim3 grid_phi(N, dim);
    int block_size = min(256, N);
    int shared_mem = block_size * 3 * sizeof(float);
    compute_phi_optimized_kernel<<<grid_phi, block_size, shared_mem>>>(
        particles.data_ptr<float>(),
        K.data_ptr<float>(),
        score.data_ptr<float>(),
        phi.data_ptr<float>(),
        inv_h,
        N, dim
    );
    
    return phi;
}

std::vector<torch::Tensor> svgd_forward_batch_cuda(torch::Tensor particles) {
    auto phi = svgd_forward_cuda(particles);
    return {phi};
}
'''

cpp_source = r'''
std::vector<torch::Tensor> svgd_forward_batch_cuda(torch::Tensor particles);
'''

svgd_module = load_inline(
    name='svgd_cuda',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['svgd_forward_batch_cuda'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    verbose=False
)


class ModelNew(nn.Module):
    """
    Stein Variational Gradient Descent (SVGD) update step.
    CUDA-optimized implementation.
    """

    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:
        # Ensure contiguous float32 input on same device
        particles = particles.contiguous().float()
        
        # Call CUDA kernel
        results = svgd_module.svgd_forward_batch_cuda(particles)
        return results[0]