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
import math
from torch.utils.cpp_extension import load_inline
cpp_src = r"""
#include <torch/extension.h>
#include <vector>
std::vector<torch::Tensor> reparam_forward(torch::Tensor mu, torch::Tensor log_var, torch::Tensor comp, torch::Tensor eps, torch::Tensor logw);
"""
cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <vector>
#include <math.h>
#define LOG_SQRT_2PI 0.9189385175704956f
__inline__ __device__ float warp_sum(float v){
    v += __shfl_down_sync(0xffffffff, v, 16);
    v += __shfl_down_sync(0xffffffff, v, 8);
    v += __shfl_down_sync(0xffffffff, v, 4);
    v += __shfl_down_sync(0xffffffff, v, 2);
    v += __shfl_down_sync(0xffffffff, v, 1);
    return v;
}
__global__ void mix4_kernel(const float* __restrict__ mu,const float* __restrict__ lv,const long long* __restrict__ comp,const float* __restrict__ eps,const float* __restrict__ logw,float* __restrict__ z,float* __restrict__ logq){
    int b=blockIdx.x, t=threadIdx.x, c=(int)comp[b]; int base=b*256+t; int sel=b*256+c*64+t;
    float lvc=lv[sel]; float zv=mu[sel]+expf(0.5f*lvc)*eps[b*64+t]; z[b*64+t]=zv;
    float m0=mu[base], l0=lv[base], d0=zv-m0; float a0=-0.5f*d0*d0*expf(-l0)-0.5f*l0-LOG_SQRT_2PI;
    float m1=mu[base+64], l1=lv[base+64], d1=zv-m1; float a1=-0.5f*d1*d1*expf(-l1)-0.5f*l1-LOG_SQRT_2PI;
    float m2=mu[base+128], l2=lv[base+128], d2=zv-m2; float a2=-0.5f*d2*d2*expf(-l2)-0.5f*l2-LOG_SQRT_2PI;
    float m3=mu[base+192], l3=lv[base+192], d3=zv-m3; float a3=-0.5f*d3*d3*expf(-l3)-0.5f*l3-LOG_SQRT_2PI;
    a0=warp_sum(a0); a1=warp_sum(a1); a2=warp_sum(a2); a3=warp_sum(a3);
    __shared__ float part[8]; int lane=t&31, w=t>>5;
    if(lane==0){part[w*4]=a0; part[w*4+1]=a1; part[w*4+2]=a2; part[w*4+3]=a3;}
    __syncthreads();
    if(t==0){
        float v0=part[0]+part[4]+logw[0]; float v1=part[1]+part[5]+logw[1]; float v2=part[2]+part[6]+logw[2]; float v3=part[3]+part[7]+logw[3];
        float mx=fmaxf(fmaxf(v0,v1),fmaxf(v2,v3));
        logq[b]=mx+logf(expf(v0-mx)+expf(v1-mx)+expf(v2-mx)+expf(v3-mx));
    }
}
__global__ void generic_kernel(const float* __restrict__ mu,const float* __restrict__ lv,const long long* __restrict__ comp,const float* __restrict__ eps,const float* __restrict__ logw,float* __restrict__ z,float* __restrict__ logq,int K,int D){
    int b=blockIdx.x,t=threadIdx.x,c=(int)comp[b]; float acc[8];
#pragma unroll
    for(int k=0;k<8;k++) acc[k]=0.f;
    for(int d=t;d<D;d+=blockDim.x){int s=(b*K+c)*D+d; float lvc=lv[s]; float zv=mu[s]+expf(.5f*lvc)*eps[b*D+d]; z[b*D+d]=zv; for(int k=0;k<K;k++){int i=(b*K+k)*D+d; float diff=zv-mu[i]; float l=lv[i]; acc[k]+=-.5f*diff*diff*expf(-l)-.5f*l-LOG_SQRT_2PI;}}
    extern __shared__ float sm[]; for(int k=0;k<K;k++) sm[k*blockDim.x+t]=acc[k]; __syncthreads(); for(int s=blockDim.x/2;s;s/=2){if(t<s){for(int k=0;k<K;k++) sm[k*blockDim.x+t]+=sm[k*blockDim.x+t+s];} __syncthreads();} if(t==0){float mx=-3.4e38f,v[8]; for(int k=0;k<K;k++){v[k]=sm[k*blockDim.x]+logw[k]; mx=fmaxf(mx,v[k]);} float se=0; for(int k=0;k<K;k++) se+=expf(v[k]-mx); logq[b]=mx+logf(se);} }
std::vector<torch::Tensor> reparam_forward(torch::Tensor mu,torch::Tensor log_var,torch::Tensor comp,torch::Tensor eps,torch::Tensor logw){int B=mu.size(0),D=eps.size(1),K=mu.dim()==3?mu.size(1):1; auto z=torch::empty({B,D},mu.options()); auto q=torch::empty({B},mu.options()); auto st=at::cuda::getCurrentCUDAStream(); if(K==4 && D==64){mix4_kernel<<<B,64,0,st>>>(mu.data_ptr<float>(),log_var.data_ptr<float>(),(const long long*)comp.data_ptr<int64_t>(),eps.data_ptr<float>(),logw.data_ptr<float>(),z.data_ptr<float>(),q.data_ptr<float>());} else {generic_kernel<<<B,64,(size_t)K*64*sizeof(float),st>>>(mu.data_ptr<float>(),log_var.data_ptr<float>(),(const long long*)comp.data_ptr<int64_t>(),eps.data_ptr<float>(),logw.data_ptr<float>(),z.data_ptr<float>(),q.data_ptr<float>(),K,D);} return {z,q};}
"""
_ext=load_inline(name='reparam_ext_v4warp_final_submit',cpp_sources=cpp_src,cuda_sources=cuda_src,functions=['reparam_forward'],extra_cuda_cflags=['-O3','--use_fast_math'],extra_cflags=['-O3'],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)); self.register_buffer('_weights',torch.full((num_components,),1.0/float(num_components))); self.register_buffer('_logw',torch.full((num_components,),-math.log(float(num_components))))
    def forward(self, mu, log_var):
        B=mu.shape[0]
        if self.num_components>1 and mu.dim()==3:
            comp=torch.multinomial(self._weights.expand(B,-1),1).squeeze(-1); eps=torch.randn((B,mu.shape[2]),device=mu.device,dtype=mu.dtype); return tuple(_ext.reparam_forward(mu,log_var,comp,eps,self._logw))
        else:
            D=mu.shape[-1]; comp=torch.zeros((B,),device=mu.device,dtype=torch.long); eps=torch.randn((B,D),device=mu.device,dtype=mu.dtype); lw=torch.zeros((1,),device=mu.device,dtype=mu.dtype); return tuple(_ext.reparam_forward(mu,log_var,comp,eps,lw))