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

_cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

__device__ float logsumexp_device(float* vals, int n) {
    float maxv = -INFINITY;
    for (int i = 0; i < n; i++) {
        if (vals[i] > maxv) maxv = vals[i];
    }
    float sum = 0.0f;
    for (int i = 0; i < n; i++) {
        sum += expf(vals[i] - maxv);
    }
    return maxv + logf(sum);
}

__global__ void importance_sampling_kernel(
    const float* __restrict__ samples,
    const float* __restrict__ mix_mus,
    const float* __restrict__ mix_stds,
    const float* __restrict__ mix_weights_log,
    const float* __restrict__ q_mu,
    const float* __restrict__ q_std_log,
    float* __restrict__ log_w_out,
    int N,
    int D,
    int K
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;

    float log_comp[8];
    for (int k = 0; k < K; k++) {
        float sum = 0.0f;
        for (int d = 0; d < D; d++) {
            float diff = samples[i * D + d] - mix_mus[k * D + d];
            float stdv = mix_stds[k * D + d];
            sum += -0.5f * (diff / stdv) * (diff / stdv) - logf(stdv);
        }
        log_comp[k] = sum + mix_weights_log[k];
    }
    float log_p = logsumexp_device(log_comp, K);

    float sum_q = 0.0f;
    for (int d = 0; d < D; d++) {
        float diff = samples[i * D + d] - q_mu[d];
        float std_log = q_std_log[d];
        sum_q += -0.5f * (diff * diff) * expf(-2.0f * std_log) - std_log;
    }
    float log_q = sum_q;

    float lw = log_p - log_q;
    log_w_out[i] = lw;
}

__global__ void normalize_and_reduce(
    const float* __restrict__ log_w,
    const float* __restrict__ samples,
    float* __restrict__ weighted_mean,
    float* __restrict__ ess_out,
    int N,
    int D
) {
    extern __shared__ float sdata[];

    float max_log = -INFINITY;
    for (int i = threadIdx.x; i < N; i += blockDim.x) {
        max_log = fmaxf(max_log, log_w[i]);
    }
    sdata[threadIdx.x] = max_log;
    __syncthreads();
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) {
            sdata[threadIdx.x] = fmaxf(sdata[threadIdx.x], sdata[threadIdx.x + s]);
        }
        __syncthreads();
    }
    max_log = sdata[0];
    __syncthreads();

    float sum_exp = 0.0f;
    float sum_w2 = 0.0f;
    for (int i = threadIdx.x; i < N; i += blockDim.x) {
        float w = expf(log_w[i] - max_log);
        sum_exp += w;
        sum_w2 += w * w;
    }
    sdata[threadIdx.x] = sum_exp;
    __syncthreads();
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) sdata[threadIdx.x] += sdata[threadIdx.x + s];
        __syncthreads();
    }
    sum_exp = sdata[0];
    sdata[threadIdx.x] = sum_w2;
    __syncthreads();
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) sdata[threadIdx.x] += sdata[threadIdx.x + s];
        __syncthreads();
    }
    sum_w2 = sdata[0];

    float log_sum = max_log + logf(sum_exp);
    if (threadIdx.x < D) {
        float wsum = 0.0f;
        for (int i = 0; i < N; i++) {
            float w = expf(log_w[i] - log_sum);
            wsum += w * samples[i * D + threadIdx.x];
        }
        weighted_mean[threadIdx.x] = wsum;
    }

    if (threadIdx.x == 0) {
        ess_out[0] = 1.0f / sum_w2;
    }
}

torch::Tensor compute_log_w(torch::Tensor samples, torch::Tensor mix_mus, torch::Tensor mix_stds, torch::Tensor mix_weights, torch::Tensor q_mu, torch::Tensor q_std) {
    int N = samples.size(0);
    int D = samples.size(1);
    int K = mix_mus.size(0);
    auto log_w = torch::empty({N}, samples.options());
    
    auto mix_weights_log = torch::log(mix_weights);
    auto q_std_log = torch::log(q_std);

    dim3 block(256);
    dim3 grid((N + 255) / 256);
    importance_sampling_kernel<<<grid, block>>>(
        samples.data_ptr<float>(),
        mix_mus.data_ptr<float>(),
        mix_stds.data_ptr<float>(),
        mix_weights_log.data_ptr<float>(),
        q_mu.data_ptr<float>(),
        q_std_log.data_ptr<float>(),
        log_w.data_ptr<float>(),
        N, D, K
    );
    return log_w;
}

std::vector<torch::Tensor> finalize(torch::Tensor log_w, torch::Tensor samples) {
    int N = samples.size(0);
    int D = samples.size(1);
    auto weighted_mean = torch::empty({D}, samples.options());
    auto ess = torch::empty({1}, samples.options());
    
    dim3 block(256);
    normalize_and_reduce<<<1, block, 256 * sizeof(float)>>>(
        log_w.data_ptr<float>(),
        samples.data_ptr<float>(),
        weighted_mean.data_ptr<float>(),
        ess.data_ptr<float>(),
        N, D
    );
    return {weighted_mean, ess.squeeze()};
}
"""

_cpp_source = """
torch::Tensor compute_log_w(torch::Tensor samples, torch::Tensor mix_mus, torch::Tensor mix_stds, torch::Tensor mix_weights, torch::Tensor q_mu, torch::Tensor q_std);
std::vector<torch::Tensor> finalize(torch::Tensor log_w, torch::Tensor samples);
"""

module = load_inline(
    name='importance_sampling',
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=['compute_log_w', 'finalize'],
    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:
        log_w = module.compute_log_w(samples, self.mix_mus, self.mix_stds, self.mix_weights, self.q_mu, self.q_std)
        return module.finalize(log_w, samples)
"""