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 compute_dist_cuda(torch::Tensor x);
torch::Tensor compute_phi_cuda(torch::Tensor x, torch::Tensor dist, torch::Tensor median);
'''

_cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be CUDA")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

__global__ void dist_kernel_generic(const float* __restrict__ x, float* __restrict__ dist, int N, int D) {
    int j = blockIdx.x * blockDim.x + threadIdx.x;
    int i = blockIdx.y * blockDim.y + threadIdx.y;
    if (i < N && j < N) {
        float s = 0.f;
        const float* xi = x + i * D;
        const float* xj = x + j * D;
        for (int d = 0; d < D; ++d) {
            float diff = xj[d] - xi[d];
            s = fmaf(diff, diff, s);
        }
        dist[i * N + j] = s;
    }
}

__global__ void dist_kernel_32_tiled(const float* __restrict__ x, float* __restrict__ dist, int N) {
    __shared__ float si[16][32];
    __shared__ float sj[16][32];
    int tx = threadIdx.x;
    int ty = threadIdx.y;
    int tid = ty * 16 + tx;
    int ib = blockIdx.y * 16;
    int jb = blockIdx.x * 16;

    for (int idx = tid; idx < 512; idx += 256) {
        int r = idx >> 5;
        int d = idx & 31;
        int gi = ib + r;
        int gj = jb + r;
        si[r][d] = (gi < N) ? x[gi * 32 + d] : 0.f;
        sj[r][d] = (gj < N) ? x[gj * 32 + d] : 0.f;
    }
    __syncthreads();

    int i = ib + ty;
    int j = jb + tx;
    if (i < N && j < N) {
        float s = 0.f;
        #pragma unroll
        for (int d = 0; d < 32; ++d) {
            float diff = sj[tx][d] - si[ty][d];
            s = fmaf(diff, diff, s);
        }
        dist[i * N + j] = s;
    }
}

__global__ void phi_kernel_32_warp(const float* __restrict__ x,
                                   const float* __restrict__ dist,
                                   const float* __restrict__ median,
                                   float* __restrict__ phi,
                                   int N) {
    __shared__ float partial[8][32];
    int i = blockIdx.x;
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;

    float med = fmaxf(*median, 1.0e-5f);
    float invh = fmaxf(logf((float)N), 1.0f) / med;
    float a = invh - 1.0f;
    float xi = x[i * 32 + lane];
    float acc = 0.f;

    for (int j = wid; j < N; j += 8) {
        float k = __expf(-0.5f * dist[i * N + j] * invh);
        float xj = x[j * 32 + lane];
        acc += k * (a * xj - invh * xi);
    }

    partial[wid][lane] = acc;
    __syncthreads();

    if (wid == 0) {
        float sum = 0.f;
        #pragma unroll
        for (int w = 0; w < 8; ++w) sum += partial[w][lane];
        phi[i * 32 + lane] = sum / (float)N;
    }
}

__global__ void phi_kernel_d64(const float* __restrict__ x,
                               const float* __restrict__ dist,
                               const float* __restrict__ median,
                               float* __restrict__ phi,
                               int N,
                               int D) {
    extern __shared__ float shf[];
    int i = blockIdx.x;
    int t = threadIdx.x;
    float med = fmaxf(*median, 1.0e-5f);
    float invh = fmaxf(logf((float)N), 1.0f) / med;
    float a = invh - 1.0f;
    float acc[64];
    for (int d = 0; d < 64; ++d) acc[d] = 0.f;
    const float* xi = x + i * D;
    for (int j = t; j < N; j += blockDim.x) {
        float k = __expf(-0.5f * dist[i * N + j] * invh);
        const float* xj = x + j * D;
        for (int d = 0; d < 64; ++d) {
            if (d < D) acc[d] += k * (a * xj[d] - invh * xi[d]);
        }
    }
    for (int d = 0; d < D; ++d) shf[t * D + d] = acc[d];
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (t < stride) {
            for (int d = 0; d < D; ++d) shf[t * D + d] += shf[(t + stride) * D + d];
        }
        __syncthreads();
    }
    if (t == 0) {
        float invN = 1.f / (float)N;
        for (int d = 0; d < D; ++d) phi[i * D + d] = shf[d] * invN;
    }
}

torch::Tensor compute_dist_cuda(torch::Tensor x) {
    CHECK_INPUT(x);
    TORCH_CHECK(x.scalar_type() == at::kFloat, "only float32 supported");
    int N = (int)x.size(0);
    int D = (int)x.size(1);
    auto dist = torch::empty({N, N}, x.options());
    cudaStream_t s = at::cuda::getCurrentCUDAStream();
    dim3 block(16, 16);
    dim3 grid((N + 15) / 16, (N + 15) / 16);
    if (D == 32) {
        dist_kernel_32_tiled<<<grid, block, 0, s>>>(x.data_ptr<float>(), dist.data_ptr<float>(), N);
    } else {
        dist_kernel_generic<<<grid, block, 0, s>>>(x.data_ptr<float>(), dist.data_ptr<float>(), N, D);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return dist;
}

torch::Tensor compute_phi_cuda(torch::Tensor x, torch::Tensor dist, torch::Tensor median) {
    CHECK_INPUT(x);
    CHECK_INPUT(dist);
    CHECK_CUDA(median);
    int N = (int)x.size(0);
    int D = (int)x.size(1);
    auto phi = torch::empty_like(x);
    cudaStream_t s = at::cuda::getCurrentCUDAStream();
    if (D == 32) {
        phi_kernel_32_warp<<<N, 256, 0, s>>>(x.data_ptr<float>(), dist.data_ptr<float>(), median.data_ptr<float>(), phi.data_ptr<float>(), N);
    } else if (D <= 64) {
        int threads = 128;
        phi_kernel_d64<<<N, threads, (size_t)threads * D * sizeof(float), s>>>(x.data_ptr<float>(), dist.data_ptr<float>(), median.data_ptr<float>(), phi.data_ptr<float>(), N, D);
    } else {
        TORCH_CHECK(false, "D>64 unsupported");
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return phi;
}
'''

_svgd_ext = load_inline(
    name='svgd_ext_final_submit',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['compute_dist_cuda', 'compute_phi_cuda'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_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:
        x = particles.contiguous()
        dist = _svgd_ext.compute_dist_cuda(x)
        med = torch.median(dist.reshape(-1)).clamp_min(1e-5)
        return _svgd_ext.compute_phi_cuda(x, dist, med)