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

# Define the CUDA kernel source code
cuda_source = '''
#include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>
#include <ATen/cuda/CUDAEvent.h>

__device__ float log_normal_pdf(float x, float mu, float sigma) {
    return -0.5 * ((x - mu) / sigma) * ((x - mu) / sigma) - __logf(sigma) - 0.5 * logf(2 * 3.14159265);
}

__global__ void reparameterize_kernel(float* z, float* log_q_z, const float* mu, const float* log_var, const float* logits, int batch_size, int num_components, int latent_dim) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= batch_size * latent_dim) return;

    int batch_idx = idx / latent_dim;
    int latent_idx = idx % latent_dim;

    if (num_components > 1) {
        // Compute softmax weights for components
        float max_logit = logits[0];
        for (int k = 1; k < num_components; ++k) {
            if (logits[k] > max_logit) max_logit = logits[k];
        }
        float sum_exp = 0.0;
        for (int k = 0; k < num_components; ++k) {
            sum_exp += expf(logits[k] - max_logit);
        }
        float log_sum_exp = logf(sum_exp) + max_logit;

        // Sample a component
        float u = (float)rand() / RAND_MAX;
        float cum_prob = 0.0;
        int comp = 0;
        for (int k = 0; k < num_components; ++k) {
            float prob = expf(logits[k] - log_sum_exp);
            cum_prob += prob;
            if (u <= cum_prob) {
                comp = k;
                break;
            }
        }

        // Reparameterize
        float mu_k = mu[batch_idx * num_components * latent_dim + comp * latent_dim + latent_idx];
        float std_k = expf(0.5 * log_var[batch_idx * num_components * latent_dim + comp * latent_dim + latent_idx]);
        float eps = (float)rand() / RAND_MAX;
        z[idx] = mu_k + std_k * eps;

        // Compute log_q_z
        float log_q_z_val = log_sum_exp;
        for (int k = 0; k < num_components; ++k) {
            float mu_val = mu[batch_idx * num_components * latent_dim + k * latent_dim + latent_idx];
            float std_val = expf(0.5 * log_var[batch_idx * num_components * latent_dim + k * latent_dim + latent_idx]);
            log_q_z_val = logaddexp(log_q_z_val, logits[k] + log_normal_pdf(z[idx], mu_val, std_val));
        }
        if (latent_idx == 0) {
            log_q_z[batch_idx] = log_q_z_val;
        }
    } else {
        float mu_val = mu[batch_idx * latent_dim + latent_idx];
        float std_val = expf(0.5 * log_var[batch_idx * latent_dim + latent_idx]);
        float eps = (float)rand() / RAND_MAX;
        z[idx] = mu_val + std_val * eps;
        if (latent_idx == 0) {
            log_q_z[batch_idx] = log_normal_pdf(z[idx], mu_val, std_val);
        } else {
            atomicAdd(&log_q_z[batch_idx], log_normal_pdf(z[idx], mu_val, std_val));
        }
    }
}

__global__ void softmax_kernel(float* weights, const float* logits, int num_components) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= num_components) return;

    // Compute softmax
    float max_logit = logits[0];
    for (int k = 1; k < num_components; ++k) {
        if (logits[k] > max_logit) max_logit = logits[k];
    }
    float sum_exp = 0.0;
    for (int k = 0; k < num_components; ++k) {
        sum_exp += expf(logits[k] - max_logit);
    }
    weights[idx] = expf(logits[idx] - max_logit) / sum_exp;
}
'''

# Compile the CUDA kernel
reparameterize_module = load_inline(
    name='reparameterize_module',
    cuda_sources=cuda_source,
    extra_cuda_cflags=['-O3'],
    functions=['reparameterize_kernel', 'softmax_kernel']
)

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))
        else:
            self.logits = None

    def forward(self, mu: torch.Tensor, log_var: torch.Tensor) -> tuple:
        batch_size = mu.shape[0]
        z = torch.empty(batch_size, self.latent_dim, device=mu.device)
        log_q_z = torch.empty(batch_size, device=mu.device)

        if self.num_components > 1:
            weights = torch.empty(self.num_components, device=mu.device)
            reparameterize_module.softmax_kernel((1,),(self.num_components,),0,None,weights.data_ptr(),self.logits.data_ptr(),self.num_components)
            reparameterize_module.reparameterize_kernel((batch_size*self.latent_dim+256)//256,256,0,None,z.data_ptr(),log_q_z.data_ptr(),mu.data_ptr(),log_var.data_ptr(),self.logits.data_ptr(),batch_size,self.num_components,self.latent_dim)
        else:
            reparameterize_module.reparameterize_kernel((batch_size*self.latent_dim+256)//256,256,0,None,z.data_ptr(),log_q_z.data_ptr(),mu.data_ptr(),log_var.data_ptr(),None,batch_size,self.num_components,self.latent_dim)
        return z, log_q_z

latent_dim = 64
num_components = 4
batch_size = 256

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

def get_init_inputs():
    return [latent_dim, num_components]