import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Reparameterization trick for differentiable sampling from a diagonal
    Gaussian.  Separates stochasticity from parameters so that gradients
    flow through mu and log_var.  Also demonstrates a "mixture of
    Gaussians" reparameterization with learnable component weights.
    """

    def __init__(self, latent_dim, num_components=1):
        super().__init__()
        self.latent_dim = latent_dim
        self.num_components = num_components
        if num_components > 1:
            self.logits = nn.Parameter(torch.zeros(num_components))

    def forward(self, mu: torch.Tensor, log_var: torch.Tensor) -> tuple:
        """
        Args:
            mu:      (B, K, D) or (B, D) – means
            log_var: (B, K, D) or (B, D) – log-variances
        Returns:
            (z, log_q_z):
                z:       (B, D) – sampled latents
                log_q_z: (B,)   – log-density of sample under q
        """
        std = torch.exp(0.5 * log_var)

        if self.num_components > 1 and mu.dim() == 3:
            weights = torch.softmax(self.logits, dim=0)  # (K,)
            # Sample component
            comp = torch.multinomial(weights.expand(mu.shape[0], -1), 1).squeeze(-1)
            batch_idx = torch.arange(mu.shape[0], device=mu.device)
            mu_k = mu[batch_idx, comp]
            std_k = std[batch_idx, comp]
            z = mu_k + std_k * torch.randn_like(mu_k)
            # log q(z) under mixture
            log_probs = -0.5 * ((z.unsqueeze(1) - mu) / std) ** 2 - log_var * 0.5 - 0.5 * torch.log(torch.tensor(2 * 3.14159265))
            log_probs = log_probs.sum(-1) + weights.log().unsqueeze(0)
            log_q_z = torch.logsumexp(log_probs, dim=1)
        else:
            if mu.dim() == 3:
                mu = mu[:, 0]
                std = std[:, 0]
                log_var = log_var[:, 0]
            z = mu + std * torch.randn_like(mu)
            log_q_z = (-0.5 * ((z - mu) / std) ** 2 - 0.5 * log_var - 0.5 * torch.log(torch.tensor(2 * 3.14159265))).sum(-1)

        return z, log_q_z


latent_dim = 64
num_components = 4
batch_size = 256


def get_inputs():
    p = popcorn_pri
    mu = torch.randn(p.jitter_int(batch_size), num_components, latent_dim)
    log_var = torch.randn(p.jitter_int(batch_size), num_components, latent_dim) * 0.5
    return [mu, log_var]


def get_init_inputs():
    return [latent_dim, num_components]

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> reparam_mix_forward(torch::Tensor mu,
                                                torch::Tensor log_var,
                                                torch::Tensor log_weights,
                                                torch::Tensor comp,
                                                torch::Tensor eps);
std::vector<torch::Tensor> reparam_diag_forward(torch::Tensor mu,
                                                 torch::Tensor log_var,
                                                 torch::Tensor eps);
"""

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

static constexpr float LOG_SQRT_2PI = 0.91893853320467274178f;

__device__ __forceinline__ float block_sum(float v, float* smem) {
    int tid = threadIdx.x;
    smem[tid] = v;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }
    return smem[0];
}

__global__ void reparam_mix_kernel(const float* __restrict__ mu,
                                   const float* __restrict__ log_var,
                                   const float* __restrict__ log_weights,
                                   const long long* __restrict__ comp,
                                   const float* __restrict__ eps,
                                   float* __restrict__ z,
                                   float* __restrict__ logq,
                                   int B, int K, int D) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    __shared__ float red[256];
    __shared__ float lps[128];

    int c = (int)comp[b];
    const int base_sel = (b * K + c) * D;
    const int zbase = b * D;

    // Sample z from the already-sampled component, using the PyTorch-generated eps.
    for (int d = tid; d < D; d += blockDim.x) {
        float lv = log_var[base_sel + d];
        float s = expf(0.5f * lv);
        z[zbase + d] = mu[base_sel + d] + s * eps[zbase + d];
    }
    __syncthreads();

    // Mixture log-density: one component at a time, block-reduced over latent dim.
    // The task uses K=4; this supports up to 128 components without dynamic shared memory.
    for (int k = 0; k < K; ++k) {
        float partial = 0.0f;
        int base = (b * K + k) * D;
        for (int d = tid; d < D; d += blockDim.x) {
            float lv = log_var[base + d];
            float stdv = expf(0.5f * lv);
            float t = (z[zbase + d] - mu[base + d]) / stdv;
            partial += -0.5f * t * t - 0.5f * lv - LOG_SQRT_2PI;
        }
        float s = block_sum(partial, red);
        if (tid == 0) lps[k] = s + log_weights[k];
        __syncthreads();
    }

    if (tid == 0) {
        float m = lps[0];
        for (int k = 1; k < K; ++k) m = fmaxf(m, lps[k]);
        float acc = 0.0f;
        for (int k = 0; k < K; ++k) acc += expf(lps[k] - m);
        logq[b] = m + logf(acc);
    }
}

__global__ void reparam_diag_kernel(const float* __restrict__ mu,
                                    const float* __restrict__ log_var,
                                    const float* __restrict__ eps,
                                    float* __restrict__ z,
                                    float* __restrict__ logq,
                                    int B, int D) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    __shared__ float red[256];
    int base = b * D;

    float partial = 0.0f;
    for (int d = tid; d < D; d += blockDim.x) {
        float lv = log_var[base + d];
        float stdv = expf(0.5f * lv);
        float e = eps[base + d];
        z[base + d] = mu[base + d] + stdv * e;
        partial += -0.5f * e * e - 0.5f * lv - LOG_SQRT_2PI;
    }
    float s = block_sum(partial, red);
    if (tid == 0) logq[b] = s;
}

std::vector<torch::Tensor> reparam_mix_forward(torch::Tensor mu,
                                                torch::Tensor log_var,
                                                torch::Tensor log_weights,
                                                torch::Tensor comp,
                                                torch::Tensor eps) {
    const int B = (int)mu.size(0);
    const int K = (int)mu.size(1);
    const int D = (int)mu.size(2);
    auto z = torch::empty({B, D}, mu.options());
    auto logq = torch::empty({B}, mu.options());

    int threads = 128;
    if (D > 128) threads = 256;
    reparam_mix_kernel<<<B, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        mu.data_ptr<float>(), log_var.data_ptr<float>(), log_weights.data_ptr<float>(),
        reinterpret_cast<const long long*>(comp.data_ptr<int64_t>()), eps.data_ptr<float>(),
        z.data_ptr<float>(), logq.data_ptr<float>(), B, K, D);
    return {z, logq};
}

std::vector<torch::Tensor> reparam_diag_forward(torch::Tensor mu,
                                                 torch::Tensor log_var,
                                                 torch::Tensor eps) {
    const int B = (int)mu.size(0);
    const int D = (int)mu.size(1);
    auto z = torch::empty({B, D}, mu.options());
    auto logq = torch::empty({B}, mu.options());
    int threads = 128;
    if (D > 128) threads = 256;
    reparam_diag_kernel<<<B, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
        mu.data_ptr<float>(), log_var.data_ptr<float>(), eps.data_ptr<float>(),
        z.data_ptr<float>(), logq.data_ptr<float>(), B, D);
    return {z, logq};
}
"""

_ext = load_inline(
    name="reparam_gauss_fused_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["reparam_mix_forward", "reparam_diag_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self, latent_dim, num_components=1):
        super().__init__()
        self.latent_dim = latent_dim
        self.num_components = num_components
        if num_components > 1:
            self.logits = nn.Parameter(torch.zeros(num_components))

    def forward(self, mu: torch.Tensor, log_var: torch.Tensor) -> tuple:
        # Mixture case used by the benchmark: keep PyTorch RNG operations so samples
        # are identical under the same seed, but fuse all density arithmetic in CUDA.
        if self.num_components > 1 and mu.dim() == 3:
            weights = torch.softmax(self.logits, dim=0)
            comp = torch.multinomial(weights.expand(mu.shape[0], -1), 1).squeeze(-1)
            eps = torch.randn((mu.shape[0], mu.shape[2]), device=mu.device, dtype=mu.dtype)
            log_weights = weights.log()
            return tuple(_ext.reparam_mix_forward(mu, log_var, log_weights, comp, eps))
        else:
            if mu.dim() == 3:
                mu2 = mu[:, 0]
                lv2 = log_var[:, 0]
                if not mu2.is_contiguous():
                    mu2 = mu2.contiguous()
                    lv2 = lv2.contiguous()
            else:
                mu2 = mu
                lv2 = log_var
            eps = torch.randn((mu2.shape[0], mu2.shape[1]), device=mu2.device, dtype=mu2.dtype)
            return tuple(_ext.reparam_diag_forward(mu2, lv2, eps))