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_runtime.h>
#include <cmath>

__device__ float logsumexp_device(float a, float b) {
    if (a == -INFINITY) return b;
    if (b == -INFINITY) return a;
    float m = fmaxf(a, b);
    return m + logf(expf(a - m) + expf(b - m));
}

__global__ void compute_log_w_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,
    float* __restrict__ log_w,
    int N, int D, int K
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;
    const float* x = samples + i * D;
    float log_p = -INFINITY;
    for (int k = 0; k < K; ++k) {
        float log_comp = mix_weights_log[k];
        for (int d = 0; d < D; ++d) {
            float diff = x[d] - mix_mus[k*D + d];
            float stdv = mix_stds[k*D + d];
            log_comp += -0.5f * diff*diff / (stdv*stdv) - logf(stdv);
        }
        log_p = logsumexp_device(log_p, log_comp);
    }
    float log_q = 0.0f;
    for (int d = 0; d < D; ++d) {
        float diff = x[d] - q_mu[d];
        float stdv = q_std[d];
        log_q += -0.5f * diff*diff / (stdv*stdv) - logf(stdv);
    }
    log_w[i] = log_p - log_q;
}

torch::Tensor compute_log_w_cuda(torch::Tensor samples, torch::Tensor mix_mus, torch::Tensor mix_stds,
                                 torch::Tensor mix_weights_log, 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());
    dim3 block(256);
    dim3 grid((N + 255) / 256);
    compute_log_w_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.data_ptr<float>(),
        log_w.data_ptr<float>(),
        N, D, K
    );
    return log_w;
}
'''

cpp_source = '''
torch::Tensor compute_log_w_cuda(torch::Tensor samples, torch::Tensor mix_mus, torch::Tensor mix_stds,
                                 torch::Tensor mix_weights_log, torch::Tensor q_mu, torch::Tensor q_std);
'''

module = load_inline(
    name='importance_sampling',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['compute_log_w_cuda'],
    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:
        mix_weights_log = self.mix_weights.log()
        log_w = module.compute_log_w_cuda(
            samples, self.mix_mus, self.mix_stds, mix_weights_log, self.q_mu, self.q_std
        )
        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