import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Self-normalized importance sampling estimate.
    Given samples from a proposal distribution q, computes importance
    weights w_i = p(x_i)/q(x_i), normalizes them, and returns both
    the weighted estimate of E_p[f(x)] and the effective sample size.

    Target p is a mixture of Gaussians; proposal q is a single Gaussian.
    """

    def __init__(self, dim, num_components=3):
        super().__init__()
        self.dim = dim
        self.num_components = num_components
        # Target mixture components
        self.register_buffer("mix_weights", torch.ones(num_components) / num_components)
        mus = torch.randn(num_components, dim) * 2.0
        self.register_buffer("mix_mus", mus)
        self.register_buffer("mix_stds", torch.ones(num_components, dim))
        # Proposal
        self.register_buffer("q_mu", torch.zeros(dim))
        self.register_buffer("q_std", torch.ones(dim) * 3.0)

    def _log_p(self, x):
        """Log density of target mixture."""
        diff = x.unsqueeze(1) - self.mix_mus.unsqueeze(0)  # (B, K, D)
        log_comp = -0.5 * (diff / self.mix_stds.unsqueeze(0)) ** 2 - self.mix_stds.log().unsqueeze(0)
        log_comp = log_comp.sum(-1)  # (B, K)
        return torch.logsumexp(log_comp + self.mix_weights.log().unsqueeze(0), dim=-1)

    def _log_q(self, x):
        """Log density of proposal."""
        return (-0.5 * ((x - self.q_mu) / self.q_std) ** 2 - self.q_std.log()).sum(-1)

    def forward(self, samples: torch.Tensor) -> tuple:
        """
        Args:
            samples: (N, dim) – samples drawn from q
        Returns:
            (weighted_mean, ess):
                weighted_mean: (dim,) – importance-weighted mean estimate of E_p[x]
                ess: scalar – effective sample size
        """
        log_w = self._log_p(samples) - self._log_q(samples)
        log_w = log_w - torch.logsumexp(log_w, dim=0)
        w = torch.exp(log_w)

        weighted_mean = (w.unsqueeze(-1) * samples).sum(dim=0)
        ess = 1.0 / (w ** 2).sum()
        return weighted_mean, ess


dim = 16
num_samples = 4096


def get_inputs():
    q_std = 3.0
    n = popcorn_pri.jitter_int(num_samples)
    return [torch.randn(n, dim) * q_std]


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>
#include <vector>

std::vector<torch::Tensor> snis_forward_cuda(torch::Tensor samples,
                                             torch::Tensor mix_weights,
                                             torch::Tensor mix_mus,
                                             torch::Tensor mix_stds,
                                             torch::Tensor q_mu,
                                             torch::Tensor q_std);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#include <cfloat>
#include <cmath>

__global__ void logw_max_d16k3_kernel(const float* __restrict__ x,
                                      const float* __restrict__ mix_w,
                                      const float* __restrict__ mus,
                                      const float* __restrict__ stds,
                                      const float* __restrict__ q_mu,
                                      const float* __restrict__ q_std,
                                      float* __restrict__ logw,
                                      float* __restrict__ max_out,
                                      int N) {
    __shared__ float smem[256];
    int tid = threadIdx.x;
    float local_max = -FLT_MAX;

    for (int i = tid; i < N; i += blockDim.x) {
        const float* xi = x + ((long long)i) * 16;
        float xv0 = xi[0],  xv1 = xi[1],  xv2 = xi[2],  xv3 = xi[3];
        float xv4 = xi[4],  xv5 = xi[5],  xv6 = xi[6],  xv7 = xi[7];
        float xv8 = xi[8],  xv9 = xi[9],  xv10 = xi[10], xv11 = xi[11];
        float xv12 = xi[12], xv13 = xi[13], xv14 = xi[14], xv15 = xi[15];

        float qlog = 0.0f;
        float d;
        d=(xv0-q_mu[0])/q_std[0];   qlog += -0.5f*d*d - logf(q_std[0]);
        d=(xv1-q_mu[1])/q_std[1];   qlog += -0.5f*d*d - logf(q_std[1]);
        d=(xv2-q_mu[2])/q_std[2];   qlog += -0.5f*d*d - logf(q_std[2]);
        d=(xv3-q_mu[3])/q_std[3];   qlog += -0.5f*d*d - logf(q_std[3]);
        d=(xv4-q_mu[4])/q_std[4];   qlog += -0.5f*d*d - logf(q_std[4]);
        d=(xv5-q_mu[5])/q_std[5];   qlog += -0.5f*d*d - logf(q_std[5]);
        d=(xv6-q_mu[6])/q_std[6];   qlog += -0.5f*d*d - logf(q_std[6]);
        d=(xv7-q_mu[7])/q_std[7];   qlog += -0.5f*d*d - logf(q_std[7]);
        d=(xv8-q_mu[8])/q_std[8];   qlog += -0.5f*d*d - logf(q_std[8]);
        d=(xv9-q_mu[9])/q_std[9];   qlog += -0.5f*d*d - logf(q_std[9]);
        d=(xv10-q_mu[10])/q_std[10]; qlog += -0.5f*d*d - logf(q_std[10]);
        d=(xv11-q_mu[11])/q_std[11]; qlog += -0.5f*d*d - logf(q_std[11]);
        d=(xv12-q_mu[12])/q_std[12]; qlog += -0.5f*d*d - logf(q_std[12]);
        d=(xv13-q_mu[13])/q_std[13]; qlog += -0.5f*d*d - logf(q_std[13]);
        d=(xv14-q_mu[14])/q_std[14]; qlog += -0.5f*d*d - logf(q_std[14]);
        d=(xv15-q_mu[15])/q_std[15]; qlog += -0.5f*d*d - logf(q_std[15]);

        float c0 = logf(mix_w[0]);
        float c1 = logf(mix_w[1]);
        float c2 = logf(mix_w[2]);
        #pragma unroll
        for (int j = 0; j < 16; ++j) {
            float xx;
            if (j==0) xx=xv0; else if (j==1) xx=xv1; else if (j==2) xx=xv2; else if (j==3) xx=xv3;
            else if (j==4) xx=xv4; else if (j==5) xx=xv5; else if (j==6) xx=xv6; else if (j==7) xx=xv7;
            else if (j==8) xx=xv8; else if (j==9) xx=xv9; else if (j==10) xx=xv10; else if (j==11) xx=xv11;
            else if (j==12) xx=xv12; else if (j==13) xx=xv13; else if (j==14) xx=xv14; else xx=xv15;
            float s0 = stds[j];
            float t = (xx - mus[j]) / s0;       c0 += -0.5f*t*t - logf(s0);
            float s1 = stds[16 + j];
            t = (xx - mus[16 + j]) / s1;        c1 += -0.5f*t*t - logf(s1);
            float s2 = stds[32 + j];
            t = (xx - mus[32 + j]) / s2;        c2 += -0.5f*t*t - logf(s2);
        }
        float m = fmaxf(c0, fmaxf(c1, c2));
        float lp = m + logf(expf(c0 - m) + expf(c1 - m) + expf(c2 - m));
        float lw = lp - qlog;
        logw[i] = lw;
        local_max = fmaxf(local_max, lw);
    }
    smem[tid] = local_max;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) smem[tid] = fmaxf(smem[tid], smem[tid + s]);
        __syncthreads();
    }
    if (tid == 0) max_out[0] = smem[0];
}

__global__ void reduce_weighted_d16_kernel(const float* __restrict__ x,
                                           const float* __restrict__ logw,
                                           const float* __restrict__ max_in,
                                           float* __restrict__ mean,
                                           float* __restrict__ ess,
                                           int N) {
    extern __shared__ float sh[];
    int tid = threadIdx.x;
    const int BD = 256;
    float m = max_in[0];
    float sw = 0.0f, sw2 = 0.0f;
    float sx[16];
    #pragma unroll
    for (int j = 0; j < 16; ++j) sx[j] = 0.0f;

    for (int i = tid; i < N; i += blockDim.x) {
        float w = expf(logw[i] - m);
        sw += w;
        sw2 += w * w;
        const float* xi = x + ((long long)i) * 16;
        #pragma unroll
        for (int j = 0; j < 16; ++j) sx[j] += w * xi[j];
    }
    sh[tid] = sw;
    sh[BD + tid] = sw2;
    #pragma unroll
    for (int j = 0; j < 16; ++j) sh[(2 + j) * BD + tid] = sx[j];
    __syncthreads();

    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) {
            sh[tid] += sh[tid + s];
            sh[BD + tid] += sh[BD + tid + s];
            #pragma unroll
            for (int j = 0; j < 16; ++j) sh[(2 + j) * BD + tid] += sh[(2 + j) * BD + tid + s];
        }
        __syncthreads();
    }
    float inv = 1.0f / sh[0];
    if (tid < 16) mean[tid] = sh[(2 + tid) * BD] * inv;
    if (tid == 0) ess[0] = (sh[0] * sh[0]) / sh[BD];
}

__global__ void logw_max_generic_kernel(const float* __restrict__ x,
                                        const float* __restrict__ mix_w,
                                        const float* __restrict__ mus,
                                        const float* __restrict__ stds,
                                        const float* __restrict__ q_mu,
                                        const float* __restrict__ q_std,
                                        float* __restrict__ logw,
                                        float* __restrict__ max_out,
                                        int N, int D, int K) {
    __shared__ float smem[256];
    int tid = threadIdx.x;
    float local_max = -FLT_MAX;
    for (int i = tid; i < N; i += blockDim.x) {
        const float* xi = x + ((long long)i) * D;
        float qlog = 0.0f;
        for (int j = 0; j < D; ++j) {
            float t = (xi[j] - q_mu[j]) / q_std[j];
            qlog += -0.5f * t * t - logf(q_std[j]);
        }
        float maxc = -FLT_MAX;
        float sumexp = 0.0f;
        for (int k = 0; k < K; ++k) {
            float c = logf(mix_w[k]);
            for (int j = 0; j < D; ++j) {
                float st = stds[k * D + j];
                float t = (xi[j] - mus[k * D + j]) / st;
                c += -0.5f * t * t - logf(st);
            }
            if (k == 0) { maxc = c; sumexp = 1.0f; }
            else {
                if (c > maxc) { sumexp = sumexp * expf(maxc - c) + 1.0f; maxc = c; }
                else sumexp += expf(c - maxc);
            }
        }
        float lw = maxc + logf(sumexp) - qlog;
        logw[i] = lw;
        local_max = fmaxf(local_max, lw);
    }
    smem[tid] = local_max;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) smem[tid] = fmaxf(smem[tid], smem[tid + s]);
        __syncthreads();
    }
    if (tid == 0) max_out[0] = smem[0];
}

__global__ void reduce_generic_kernel(const float* __restrict__ x,
                                      const float* __restrict__ logw,
                                      const float* __restrict__ max_in,
                                      float* __restrict__ mean,
                                      float* __restrict__ ess,
                                      int N, int D) {
    __shared__ float red[256];
    int tid = threadIdx.x;
    float m = max_in[0];
    // First reduce scalar sums.
    float sw = 0.0f, sw2 = 0.0f;
    for (int i = tid; i < N; i += blockDim.x) {
        float w = expf(logw[i] - m);
        sw += w; sw2 += w*w;
    }
    red[tid] = sw; __syncthreads();
    for (int s=blockDim.x>>1; s>0; s>>=1) { if (tid<s) red[tid]+=red[tid+s]; __syncthreads(); }
    float total = red[0];
    red[tid] = sw2; __syncthreads();
    for (int s=blockDim.x>>1; s>0; s>>=1) { if (tid<s) red[tid]+=red[tid+s]; __syncthreads(); }
    if (tid == 0) ess[0] = total*total/red[0];
    for (int d = 0; d < D; ++d) {
        float sx = 0.0f;
        for (int i = tid; i < N; i += blockDim.x) sx += expf(logw[i] - m) * x[(long long)i * D + d];
        red[tid] = sx; __syncthreads();
        for (int s=blockDim.x>>1; s>0; s>>=1) { if (tid<s) red[tid]+=red[tid+s]; __syncthreads(); }
        if (tid == 0) mean[d] = red[0] / total;
        __syncthreads();
    }
}

std::vector<torch::Tensor> snis_forward_cuda(torch::Tensor samples,
                                             torch::Tensor mix_weights,
                                             torch::Tensor mix_mus,
                                             torch::Tensor mix_stds,
                                             torch::Tensor q_mu,
                                             torch::Tensor q_std) {
    TORCH_CHECK(samples.is_cuda(), "samples must be CUDA");
    TORCH_CHECK(samples.scalar_type() == torch::kFloat32, "only float32 supported");
    TORCH_CHECK(samples.is_contiguous(), "samples must be contiguous");
    int N = (int)samples.size(0);
    int D = (int)samples.size(1);
    int K = (int)mix_weights.size(0);
    auto opts = samples.options();
    auto logw = torch::empty({N}, opts);
    auto maxv = torch::empty({1}, opts);
    auto mean = torch::empty({D}, opts);
    auto ess = torch::empty({}, opts);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int threads = 256;
    if (D == 16 && K == 3) {
        logw_max_d16k3_kernel<<<1, threads, 0, stream>>>(
            samples.data_ptr<float>(), mix_weights.data_ptr<float>(), mix_mus.data_ptr<float>(),
            mix_stds.data_ptr<float>(), q_mu.data_ptr<float>(), q_std.data_ptr<float>(),
            logw.data_ptr<float>(), maxv.data_ptr<float>(), N);
        reduce_weighted_d16_kernel<<<1, threads, 18 * threads * sizeof(float), stream>>>(
            samples.data_ptr<float>(), logw.data_ptr<float>(), maxv.data_ptr<float>(),
            mean.data_ptr<float>(), ess.data_ptr<float>(), N);
    } else {
        logw_max_generic_kernel<<<1, threads, 0, stream>>>(
            samples.data_ptr<float>(), mix_weights.data_ptr<float>(), mix_mus.data_ptr<float>(),
            mix_stds.data_ptr<float>(), q_mu.data_ptr<float>(), q_std.data_ptr<float>(),
            logw.data_ptr<float>(), maxv.data_ptr<float>(), N, D, K);
        reduce_generic_kernel<<<1, threads, 0, stream>>>(
            samples.data_ptr<float>(), logw.data_ptr<float>(), maxv.data_ptr<float>(),
            mean.data_ptr<float>(), ess.data_ptr<float>(), N, D);
    }
    return {mean, ess};
}
"""

_ext = load_inline(
    name="snis_importance_ext_v1",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["snis_forward_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, dim, num_components=3):
        super().__init__()
        self.dim = dim
        self.num_components = num_components
        self.register_buffer("mix_weights", torch.ones(num_components) / num_components)
        mus = torch.randn(num_components, dim) * 2.0
        self.register_buffer("mix_mus", mus)
        self.register_buffer("mix_stds", torch.ones(num_components, dim))
        self.register_buffer("q_mu", torch.zeros(dim))
        self.register_buffer("q_std", torch.ones(dim) * 3.0)

    def forward(self, samples: torch.Tensor) -> tuple:
        out = _ext.snis_forward_cuda(samples.contiguous(), self.mix_weights.contiguous(),
                                     self.mix_mus.contiguous(), self.mix_stds.contiguous(),
                                     self.q_mu.contiguous(), self.q_std.contiguous())
        return (out[0], out[1])