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 = r'''
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

__global__ void importance_sampling_kernel(
    const float* samples,
    const float* mix_mus,
    const float* mix_stds,
    const float* log_mix_weights,
    const float* q_mu,
    const float* q_std,
    float* log_weights,
    int num_samples, int dim, int num_components
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (idx < num_samples) {
        float max_log_comp = -1e30f;
        float sum_exp = 0.0f;
        
        for (int k = 0; k < num_components; k++) {
            float log_comp = log_mix_weights[k];
            for (int d = 0; d < dim; d++) {
                float diff = samples[idx * dim + d] - mix_mus[k * dim + d];
                float std = mix_stds[k * dim + d];
                float z = diff / std;
                log_comp += -0.5f * z * z - logf(std);
            }
            
            if (log_comp > max_log_comp) {
                sum_exp = expf(max_log_comp - log_comp) * sum_exp + 1.0f;
                max_log_comp = log_comp;
            } else {
                sum_exp += expf(log_comp - max_log_comp);
            }
        }
        
        float log_p = max_log_comp + logf(sum_exp + 1e-30f);
        
        float log_q = 0.0f;
        for (int d = 0; d < dim; d++) {
            float diff = samples[idx * dim + d] - q_mu[d];
            float std = q_std[d];
            float z = diff / std;
            log_q += -0.5f * z * z - logf(std);
        }
        
        log_weights[idx] = log_p - log_q;
    }
}

__global__ void logsumexp_kernel(
    const float* log_weights,
    float* result,
    int num_samples
) {
    __shared__ float max_val;
    
    int tid = threadIdx.x;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    float local_max = -1e30f;
    
    for (int i = idx; i < num_samples; i += blockDim.x) {
        local_max = fmaxf(local_max, log_weights[i]);
    }
    
    if (tid == 0) max_val = -1e30f;
    __syncthreads();
    
    for (int offset = 16; offset > 0; offset >>= 1) {
        local_max = fmaxf(local_max, __shfl_down_sync(0xffffffff, local_max, offset));
    }
    
    if (tid == 0) max_val = local_max;
    __syncthreads();
    
    float global_max = max_val;
    
    float local_sum = 0.0f;
    for (int i = idx; i < num_samples; i += blockDim.x) {
        local_sum += expf(log_weights[i] - global_max);
    }
    
    for (int offset = 16; offset > 0; offset >>= 1) {
        local_sum += __shfl_down_sync(0xffffffff, local_sum, offset);
    }
    
    if (tid == 0) {
        result[0] = global_max + logf(local_sum + 1e-30f);
    }
}

__global__ void normalize_and_reduce_kernel(
    const float* log_weights,
    const float* samples,
    float* weighted_mean,
    float* ess_val,
    float log_sum_weights,
    int num_samples, int dim
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (idx < num_samples) {
        float w = expf(log_weights[idx] - log_sum_weights);
        float w_sq = w * w;
        
        for (int d = 0; d < dim; d++) {
            atomicAdd(&weighted_mean[d], w * samples[idx * dim + d]);
        }
        atomicAdd(ess_val, w_sq);
    }
}

void run_importance_sampling_kernel(
    const float* samples, const float* mix_mus, const float* mix_stds,
    const float* log_mix_weights, const float* q_mu, const float* q_std,
    float* log_weights, int num_samples, int dim, int num_components
) {
    const int block_size = 256;
    const int num_blocks = (num_samples + block_size - 1) / block_size;
    importance_sampling_kernel<<<num_blocks, block_size>>>(
        samples, mix_mus, mix_stds, log_mix_weights, q_mu, q_std, log_weights,
        num_samples, dim, num_components
    );
}

void run_logsumexp_kernel(const float* log_weights, float* result, int num_samples) {
    const int block_size = 256;
    logsumexp_kernel<<<1, block_size>>>(log_weights, result, num_samples);
}

void run_normalize_and_reduce_kernel(
    const float* log_weights, const float* samples, float* weighted_mean,
    float* ess_val, float log_sum_weights, int num_samples, int dim
) {
    const int block_size = 256;
    const int num_blocks = (num_samples + block_size - 1) / block_size;
    normalize_and_reduce_kernel<<<num_blocks, block_size>>>(
        log_weights, samples, weighted_mean, ess_val, log_sum_weights, num_samples, dim
    );
}
'''

cpp_source = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>

void run_importance_sampling_kernel(
    const float* samples, const float* mix_mus, const float* mix_stds,
    const float* log_mix_weights, const float* q_mu, const float* q_std,
    float* log_weights, int num_samples, int dim, int num_components);

void run_logsumexp_kernel(const float* log_weights, float* result, int num_samples);

void run_normalize_and_reduce_kernel(
    const float* log_weights, const float* samples, float* weighted_mean,
    float* ess_val, float log_sum_weights, int num_samples, int dim);

torch::Tensor importance_sampling_forward(
    torch::Tensor samples,
    torch::Tensor mix_mus,
    torch::Tensor mix_stds,
    torch::Tensor mix_weights,
    torch::Tensor q_mu,
    torch::Tensor q_std
) {
    int num_samples = samples.size(0);
    int dim = samples.size(1);
    int num_components = mix_weights.size(0);
    
    auto log_weights = torch::empty({num_samples}, samples.options());
    auto log_mix_weights = torch::log(mix_weights);
    auto weighted_mean = torch::zeros({dim}, samples.options());
    auto ess_val = torch::zeros({1}, samples.options());
    auto log_sum_weights = torch::empty({1}, samples.options());
    
    run_importance_sampling_kernel(
        samples.data_ptr<float>(),
        mix_mus.data_ptr<float>(),
        mix_stds.data_ptr<float>(),
        log_mix_weights.data_ptr<float>(),
        q_mu.data_ptr<float>(),
        q_std.data_ptr<float>(),
        log_weights.data_ptr<float>(),
        num_samples, dim, num_components
    );
    
    run_logsumexp_kernel(
        log_weights.data_ptr<float>(),
        log_sum_weights.data_ptr<float>(),
        num_samples
    );
    
    float log_sum;
    cudaMemcpy(&log_sum, log_sum_weights.data_ptr<float>(), sizeof(float), cudaMemcpyDeviceToHost);
    
    run_normalize_and_reduce_kernel(
        log_weights.data_ptr<float>(),
        samples.data_ptr<float>(),
        weighted_mean.data_ptr<float>(),
        ess_val.data_ptr<float>(),
        log_sum,
        num_samples, dim
    );
    
    float ess_sum;
    cudaMemcpy(&ess_sum, ess_val.data_ptr<float>(), sizeof(float), cudaMemcpyDeviceToHost);
    ess_sum = 1.0f / ess_sum;
    cudaMemcpy(ess_val.data_ptr<float>(), &ess_sum, sizeof(float), cudaMemcpyHostToDevice);
    
    return torch::cat({weighted_mean, ess_val});
}
'''

module = load_inline(
    name="importance_sampling_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["importance_sampling_forward"],
    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:
        result = module.importance_sampling_forward(
            samples, self.mix_mus, self.mix_stds, self.mix_weights, self.q_mu, self.q_std
        )
        weighted_mean = result[:self.dim]
        ess = result[self.dim]
        return weighted_mean, ess