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
import math

cuda_source = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <curand_kernel.h>
#include <cmath>
#include <limits>

// Constants
static constexpr float LOG_2PI = 1.8378770664093453f;  // log(2*pi)
static constexpr float PI = 3.14159265358979323846f;

// Warp-level reduction for sum
__inline__ __device__ float warp_reduce_sum(float val) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xffffffff, val, offset);
    }
    return val;
}

// Warp-level reduction for max
__inline__ __device__ float warp_reduce_max(float val) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
    }
    return val;
}

// Single component kernel (for 2D input)
__global__ void single_component_2d_kernel(
    const float* __restrict__ mu,
    const float* __restrict__ log_var,
    float* __restrict__ z,
    float* __restrict__ log_q_z,
    const int batch_size,
    const int latent_dim,
    const uint64_t seed
) {
    const int b = blockIdx.x;
    const int tid = threadIdx.x;
    const int lane_id = tid & 31;
    const int warp_id = tid >> 5;
    const int warps_per_block = blockDim.x >> 5;
    
    if (b >= batch_size) return;
    
    // Each thread handles multiple latent dimensions
    float local_sum = 0.0f;
    
    // Initialize random state per thread
    curandState state;
    curand_init(seed, b * blockDim.x + tid, 0, &state);
    
    for (int d = tid; d < latent_dim; d += blockDim.x) {
        const int idx = b * latent_dim + d;
        const float mu_val = mu[idx];
        const float log_var_val = log_var[idx];
        const float std_val = expf(0.5f * log_var_val);
        
        const float noise = curand_normal(&state);
        const float z_val = mu_val + std_val * noise;
        z[idx] = z_val;
        
        // log_p contribution: -0.5 * ((z-mu)/std)^2 - 0.5*log_var - 0.5*log(2*pi)
        // = -0.5 * noise^2 - 0.5*log_var - 0.5*log(2*pi)
        const float log_p = -0.5f * noise * noise - 0.5f * log_var_val - 0.5f * LOG_2PI;
        local_sum += log_p;
    }
    
    // Warp-level reduction
    local_sum = warp_reduce_sum(local_sum);
    
    // First thread in each warp writes to shared memory
    __shared__ float s_partial[32];  // Max 32 warps per block
    if (lane_id == 0) {
        s_partial[warp_id] = local_sum;
    }
    __syncthreads();
    
    // Final reduction by first warp
    if (warp_id == 0) {
        float sum = (tid < warps_per_block) ? s_partial[tid] : 0.0f;
        sum = warp_reduce_sum(sum);
        if (tid == 0) {
            log_q_z[b] = sum;
        }
    }
}

// Mixture of Gaussians kernel - step 1: sample z
__global__ void mixture_sample_kernel(
    const float* __restrict__ mu,
    const float* __restrict__ log_var,
    const int* __restrict__ components,
    float* __restrict__ z,
    const int batch_size,
    const int num_components,
    const int latent_dim,
    const uint64_t seed
) {
    const int b = blockIdx.x;
    const int tid = threadIdx.x;
    
    if (b >= batch_size) return;
    
    const int comp = components[b];
    
    // Initialize random state
    curandState state;
    curand_init(seed, b * blockDim.x + tid, 0, &state);
    
    for (int d = tid; d < latent_dim; d += blockDim.x) {
        const int idx_mu = b * num_components * latent_dim + comp * latent_dim + d;
        const int idx_z = b * latent_dim + d;
        
        const float mu_val = mu[idx_mu];
        const float log_var_val = log_var[idx_mu];
        const float std_val = expf(0.5f * log_var_val);
        
        const float noise = curand_normal(&state);
        z[idx_z] = mu_val + std_val * noise;
    }
}

// Mixture of Gaussians kernel - step 2: compute log_q_z with logsumexp
__global__ void mixture_logprob_kernel(
    const float* __restrict__ mu,
    const float* __restrict__ log_var,
    const float* __restrict__ z,
    const float* __restrict__ log_weights,
    float* __restrict__ log_q_z,
    const int batch_size,
    const int num_components,
    const int latent_dim
) {
    const int b = blockIdx.x;
    const int tid = threadIdx.x;
    
    if (b >= batch_size) return;
    
    // Each thread computes log probability for one component
    // Then we do logsumexp across components
    extern __shared__ float s_log_probs[];  // size: num_components
    
    float max_log = -INFINITY;
    
    for (int k = tid; k < num_components; k += blockDim.x) {
        float sum_log_p = log_weights[k];  // Start with log(weight)
        
        for (int d = 0; d < latent_dim; d++) {
            const int idx_mu = b * num_components * latent_dim + k * latent_dim + d;
            const int idx_z = b * latent_dim + d;
            
            const float z_val = z[idx_z];
            const float mu_val = mu[idx_mu];
            const float log_var_val = log_var[idx_mu];
            const float std_val = expf(0.5f * log_var_val);
            
            // log N(z | mu, std) = -0.5 * ((z-mu)/std)^2 - 0.5*log_var - 0.5*log(2*pi)
            const float diff = (z_val - mu_val) / std_val;
            sum_log_p += -0.5f * diff * diff - 0.5f * log_var_val - 0.5f * LOG_2PI;
        }
        
        s_log_probs[k] = sum_log_p;
        max_log = fmaxf(max_log, sum_log_p);
    }
    __syncthreads();
    
    // Reduce max_log across threads
    __shared__ float s_max;
    if (tid == 0) {
        s_max = s_log_probs[0];
        for (int k = 1; k < num_components; k++) {
            s_max = fmaxf(s_max, s_log_probs[k]);
        }
    }
    __syncthreads();
    
    // Compute logsumexp
    const float max_val = s_max;
    float sum_exp = 0.0f;
    
    for (int k = tid; k < num_components; k += blockDim.x) {
        sum_exp += expf(s_log_probs[k] - max_val);
    }
    
    // Warp-level reduction
    sum_exp = warp_reduce_sum(sum_exp);
    
    __shared__ float s_sum_exp;
    if (tid == 0) {
        s_sum_exp = sum_exp;
    }
    __syncthreads();
    
    if (tid == 0) {
        log_q_z[b] = max_val + logf(s_sum_exp);
    }
}

// Host functions
std::tuple<torch::Tensor, torch::Tensor> single_component_forward(
    torch::Tensor mu, torch::Tensor log_var,
    int64_t latent_dim, uint64_t seed
) {
    const int batch_size = mu.size(0);
    
    auto z = torch::empty({batch_size, latent_dim}, mu.options());
    auto log_q_z = torch::empty({batch_size}, mu.options());
    
    const int threads = 256;
    const int blocks = batch_size;
    
    single_component_2d_kernel<<<blocks, threads>>>(
        mu.data_ptr<float>(),
        log_var.data_ptr<float>(),
        z.data_ptr<float>(),
        log_q_z.data_ptr<float>(),
        batch_size,
        latent_dim,
        seed
    );
    
    return std::make_tuple(z, log_q_z);
}

std::tuple<torch::Tensor, torch::Tensor> mixture_forward(
    torch::Tensor mu, torch::Tensor log_var,
    torch::Tensor weights,
    int64_t latent_dim, int64_t num_components,
    uint64_t seed
) {
    const int batch_size = mu.size(0);
    
    // Sample components using multinomial (PyTorch function)
    auto weights_expanded = weights.expand({batch_size, -1});
    auto components = torch::multinomial(weights_expanded, 1).squeeze(-1);
    
    auto z = torch::empty({batch_size, latent_dim}, mu.options());
    auto log_q_z = torch::empty({batch_size}, mu.options());
    auto log_weights = torch::log(weights);
    
    const int threads = 256;
    const int blocks = batch_size;
    
    // Step 1: Sample z
    mixture_sample_kernel<<<blocks, threads>>>(
        mu.data_ptr<float>(),
        log_var.data_ptr<float>(),
        components.data_ptr<int>(),
        z.data_ptr<float>(),
        batch_size,
        num_components,
        latent_dim,
        seed
    );
    
    // Step 2: Compute log_q_z
    size_t shared_mem = num_components * sizeof(float);
    mixture_logprob_kernel<<<blocks, threads, shared_mem>>>(
        mu.data_ptr<float>(),
        log_var.data_ptr<float>(),
        z.data_ptr<float>(),
        log_weights.data_ptr<float>(),
        log_q_z.data_ptr<float>(),
        batch_size,
        num_components,
        latent_dim
    );
    
    return std::make_tuple(z, log_q_z);
}

std::tuple<torch::Tensor, torch::Tensor> forward_3d_single(
    torch::Tensor mu, torch::Tensor log_var,
    int64_t latent_dim, uint64_t seed
) {
    // 3D input but single component - take first component
    const int batch_size = mu.size(0);
    
    auto mu_2d = mu.select(1, 0).contiguous();
    auto log_var_2d = log_var.select(1, 0).contiguous();
    
    return single_component_forward(mu_2d, log_var_2d, latent_dim, seed);
}
"""

cpp_source = """
std::tuple<torch::Tensor, torch::Tensor> single_component_forward(
    torch::Tensor mu, torch::Tensor log_var,
    int64_t latent_dim, uint64_t seed
);

std::tuple<torch::Tensor, torch::Tensor> mixture_forward(
    torch::Tensor mu, torch::Tensor log_var,
    torch::Tensor weights,
    int64_t latent_dim, int64_t num_components,
    uint64_t seed
);

std::tuple<torch::Tensor, torch::Tensor> forward_3d_single(
    torch::Tensor mu, torch::Tensor log_var,
    int64_t latent_dim, uint64_t seed
);
"""

# Compile the extension
module = load_inline(
    name='reparam_cuda',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['single_component_forward', 'mixture_forward', 'forward_3d_single'],
    extra_cuda_cflags=['-O3', '--use_fast_math', '-gencode=arch=compute_80,code=sm_80'],
    verbose=False
)

import os
import time

class ModelNew(nn.Module):
    """
    Reparameterization trick for differentiable sampling from a diagonal
    Gaussian. Uses custom CUDA kernels for faster computation.
    """
    
    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))
        
        # Seed for reproducibility (can be changed)
        self._seed_counter = 0
    
    def _get_seed(self):
        # Use a combination of time and counter for different seeds
        self._seed_counter += 1
        return int(time.time() * 1e6 + self._seed_counter) % (2**63)
    
    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
        """
        # Ensure contiguous float32 tensors
        mu = mu.contiguous().float()
        log_var = log_var.contiguous().float()
        
        seed = self._get_seed()
        
        if self.num_components > 1 and mu.dim() == 3:
            # Mixture of Gaussians case
            weights = torch.softmax(self.logits, dim=0)
            z, log_q_z = module.mixture_forward(
                mu, log_var, weights,
                self.latent_dim, self.num_components,
                seed
            )
        else:
            # Single component case
            if mu.dim() == 3:
                # Take first component
                mu = mu[:, 0].contiguous()
                log_var = log_var[:, 0].contiguous()
            
            z, log_q_z = module.single_component_forward(
                mu, log_var, self.latent_dim, seed
            )
        
        return z, log_q_z