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

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

__global__ void reparam_mixture_kernel(
    const float* __restrict__ mu,
    const float* __restrict__ log_var,
    const float* __restrict__ weights,
    const int64_t* __restrict__ comp,
    const float* __restrict__ randn,
    float* __restrict__ z,
    float* __restrict__ log_q_z,
    int B, int K, int D, float log_2pi
) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= B) return;

    float w_log[16];
    for(int k=0; k<K; k++) w_log[k] = logf(weights[k]);

    int c = comp[b];
    for(int d=0; d<D; d++) {
        int idx_k = (b * K + c) * D + d;
        float muk = mu[idx_k];
        float stdk = expf(0.5f * log_var[idx_k]);
        float r = randn[b * D + d];
        z[b * D + d] = muk + stdk * r;
    }

    float max_lp = -INFINITY;
    float lps[16];
    for(int k=0; k<K; k++) {
        float lp = 0.0f;
        for(int d=0; d<D; d++) {
            int idx = (b * K + k) * D + d;
            float m = mu[idx];
            float lv = log_var[idx];
            float s = expf(0.5f * lv);
            float diff = (z[b*D + d] - m) / s;
            lp += -0.5f * diff * diff - 0.5f * lv - 0.5f * log_2pi;
        }
        lp += w_log[k];
        lps[k] = lp;
        if(lp > max_lp) max_lp = lp;
    }
    float sum_exp = 0.0f;
    for(int k=0; k<K; k++) {
        sum_exp += expf(lps[k] - max_lp);
    }
    log_q_z[b] = max_lp + logf(sum_exp);
}

std::tuple<torch::Tensor, torch::Tensor> reparam_mixture_cuda(
    torch::Tensor mu,
    torch::Tensor log_var,
    torch::Tensor weights,
    torch::Tensor comp,
    torch::Tensor randn,
    int num_components
) {
    int B = mu.size(0);
    int K = num_components;
    int D = mu.size(2);
    auto z = torch::empty({B, D}, mu.options());
    auto log_q_z = torch::empty({B}, mu.options());
    float log_2pi = logf(2.0f * 3.14159265f);

    int threads = 256;
    int blocks = (B + threads - 1) / threads;
    reparam_mixture_kernel<<<blocks, threads>>>(
        mu.data_ptr<float>(),
        log_var.data_ptr<float>(),
        weights.data_ptr<float>(),
        comp.data_ptr<int64_t>(),
        randn.data_ptr<float>(),
        z.data_ptr<float>(),
        log_q_z.data_ptr<float>(),
        B, K, D, log_2pi
    );
    return std::make_tuple(z, log_q_z);
}
"""

cpp_source = """
std::tuple<torch::Tensor, torch::Tensor> reparam_mixture_cuda(
    torch::Tensor mu, torch::Tensor log_var, torch::Tensor weights,
    torch::Tensor comp, torch::Tensor randn, int num_components);
"""

module = load_inline(
    name="reparam_mixture",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["reparam_mixture_cuda"],
    verbose=False,
    extra_cuda_cflags=['-O3']
)

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:
        if self.num_components > 1 and mu.dim() == 3:
            weights = torch.softmax(self.logits, dim=0)
            B = mu.shape[0]
            comp = torch.multinomial(weights.expand(B, -1), 1).squeeze(-1)
            randn = torch.randn(B, mu.shape[2], device=mu.device, dtype=mu.dtype)
            z, log_q_z = module.reparam_mixture_cuda(mu, log_var, weights, comp, randn, self.num_components)
            return z, log_q_z
        else:
            std = torch.exp(0.5 * log_var)
            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, device=mu.device, dtype=mu.dtype))).sum(-1)
            return z, log_q_z