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);
std::vector<torch::Tensor> snis_forward(torch::Tensor samples, torch::Tensor mix_weights,
                                        torch::Tensor mix_mus, torch::Tensor mix_stds,
                                        torch::Tensor q_mu, torch::Tensor q_std) {
    return snis_forward_cuda(samples, mix_weights, mix_mus, mix_stds, q_mu, q_std);
}
"""

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

__device__ __forceinline__ float compute_logw_16_3(const float* __restrict__ x,
                                                    const float* __restrict__ mix_weights,
                                                    const float* __restrict__ mu) {
    float sx2 = 0.0f;
    float d0 = 0.0f, d1 = 0.0f, d2 = 0.0f;
    #pragma unroll
    for (int d = 0; d < 16; ++d) {
        float xd = x[d];
        sx2 += xd * xd;
        float a0 = xd - mu[d];
        float a1 = xd - mu[16 + d];
        float a2 = xd - mu[32 + d];
        d0 += a0 * a0;
        d1 += a1 * a1;
        d2 += a2 * a2;
    }
    float lw = logf(mix_weights[0]);
    float c0 = -0.5f * d0 + lw;
    float c1 = -0.5f * d1 + lw;
    float c2 = -0.5f * d2 + lw;
    float m = fmaxf(c0, fmaxf(c1, c2));
    float lse = m + logf(expf(c0 - m) + expf(c1 - m) + expf(c2 - m));
    return lse + 0.05555555555555556f * sx2 + 17.577796618689757f;
}

__global__ void logw16_kernel(const float* __restrict__ samples,
                              const float* __restrict__ mix_weights,
                              const float* __restrict__ mix_mus,
                              int N,
                              float* __restrict__ logw,
                              float* __restrict__ partial_max) {
    __shared__ float s[256];
    int tid = threadIdx.x;
    int i = blockIdx.x * 256 + tid;
    float v = -INFINITY;
    if (i < N) {
        v = compute_logw_16_3(samples + ((long long)i << 4), mix_weights, mix_mus);
        logw[i] = v;
    }
    s[tid] = v;
    __syncthreads();
    #pragma unroll
    for (int stride = 128; stride > 0; stride >>= 1) {
        if (tid < stride) s[tid] = fmaxf(s[tid], s[tid + stride]);
        __syncthreads();
    }
    if (tid == 0) partial_max[blockIdx.x] = s[0];
}

__global__ void reduce16_kernel(const float* __restrict__ samples,
                                const float* __restrict__ logw,
                                const float* __restrict__ partial_max,
                                int N, int num_blocks,
                                float* __restrict__ out_mean,
                                float* __restrict__ out_ess) {
    extern __shared__ float smem[];
    int tid = threadIdx.x;
    float* s_max = smem;
    float* s_sum = smem + 256;
    float* s_sum2 = smem + 512;
    float* s_num = smem + 768;

    float lm = -INFINITY;
    for (int j = tid; j < num_blocks; j += 256) lm = fmaxf(lm, partial_max[j]);
    s_max[tid] = lm;
    __syncthreads();
    #pragma unroll
    for (int stride = 128; stride > 0; stride >>= 1) {
        if (tid < stride) s_max[tid] = fmaxf(s_max[tid], s_max[tid + stride]);
        __syncthreads();
    }
    float maxlw = s_max[0];

    float lsum = 0.0f, lsum2 = 0.0f;
    float n0=0.0f,n1=0.0f,n2=0.0f,n3=0.0f,n4=0.0f,n5=0.0f,n6=0.0f,n7=0.0f;
    float n8=0.0f,n9=0.0f,n10=0.0f,n11=0.0f,n12=0.0f,n13=0.0f,n14=0.0f,n15=0.0f;
    for (int i = tid; i < N; i += 256) {
        float a = expf(logw[i] - maxlw);
        lsum += a;
        lsum2 += a * a;
        const float* x = samples + ((long long)i << 4);
        n0 += a*x[0]; n1 += a*x[1]; n2 += a*x[2]; n3 += a*x[3];
        n4 += a*x[4]; n5 += a*x[5]; n6 += a*x[6]; n7 += a*x[7];
        n8 += a*x[8]; n9 += a*x[9]; n10 += a*x[10]; n11 += a*x[11];
        n12 += a*x[12]; n13 += a*x[13]; n14 += a*x[14]; n15 += a*x[15];
    }
    s_sum[tid] = lsum;
    s_sum2[tid] = lsum2;
    s_num[0*256+tid]=n0; s_num[1*256+tid]=n1; s_num[2*256+tid]=n2; s_num[3*256+tid]=n3;
    s_num[4*256+tid]=n4; s_num[5*256+tid]=n5; s_num[6*256+tid]=n6; s_num[7*256+tid]=n7;
    s_num[8*256+tid]=n8; s_num[9*256+tid]=n9; s_num[10*256+tid]=n10; s_num[11*256+tid]=n11;
    s_num[12*256+tid]=n12; s_num[13*256+tid]=n13; s_num[14*256+tid]=n14; s_num[15*256+tid]=n15;
    __syncthreads();

    #pragma unroll
    for (int stride = 128; stride > 0; stride >>= 1) {
        if (tid < stride) {
            s_sum[tid] += s_sum[tid + stride];
            s_sum2[tid] += s_sum2[tid + stride];
            #pragma unroll
            for (int d = 0; d < 16; ++d) s_num[d*256 + tid] += s_num[d*256 + tid + stride];
        }
        __syncthreads();
    }
    if (tid == 0) {
        float inv = 1.0f / s_sum[0];
        #pragma unroll
        for (int d = 0; d < 16; ++d) out_mean[d] = s_num[d*256] * inv;
        out_ess[0] = s_sum[0] * s_sum[0] / s_sum2[0];
    }
}

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) {
    samples = samples.contiguous();
    int N = (int)samples.size(0);
    auto mean = torch::empty({16}, samples.options());
    auto ess = torch::empty({}, samples.options());
    auto logw = torch::empty({N}, samples.options());
    const int threads = 256;
    int blocks = (N + threads - 1) / threads;
    auto partial = torch::empty({blocks}, samples.options());
    logw16_kernel<<<blocks, threads>>>(samples.data_ptr<float>(), mix_weights.data_ptr<float>(), mix_mus.data_ptr<float>(),
                                       N, logw.data_ptr<float>(), partial.data_ptr<float>());
    reduce16_kernel<<<1, threads, 19 * threads * sizeof(float)>>>(samples.data_ptr<float>(), logw.data_ptr<float>(),
                                                                  partial.data_ptr<float>(), N, blocks,
                                                                  mean.data_ptr<float>(), ess.data_ptr<float>());
    return {mean, ess};
}
"""

_ext = load_inline(name="snis_ext_v3", cpp_sources=_cpp_src, cuda_sources=_cuda_src,
                   functions=["snis_forward"], verbose=False, extra_cuda_cflags=["-O3", "--use_fast_math"])

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(samples, self.mix_weights, self.mix_mus, self.mix_stds, self.q_mu, self.q_std)
        return (out[0], out[1])