import torch
import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    """
    Evidence Lower Bound (ELBO) computation for a Variational Autoencoder.
    Combines a reconstruction term (binary cross-entropy) with the
    KL divergence between the approximate posterior q(z|x) = N(mu, sigma^2)
    and the prior p(z) = N(0, I).  The core loss function in VAE training.
    """

    def __init__(self, input_dim, hidden_dim, latent_dim):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
        )
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)

        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim),
            nn.Sigmoid(),
        )

    def forward(self, x: torch.Tensor) -> tuple:
        """
        Args:
            x: (B, input_dim) – input data in [0, 1]
        Returns:
            (elbo, recon_loss, kl_div):
                elbo:       (B,)
                recon_loss: (B,)
                kl_div:     (B,)
        """
        h = self.encoder(x)
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)

        std = torch.exp(0.5 * logvar)
        z = mu + std * torch.randn_like(std)

        x_recon = self.decoder(z)
        recon = F.binary_cross_entropy(x_recon, x, reduction="none").sum(-1)
        kl = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()).sum(-1)
        elbo = -(recon + kl)
        return elbo, recon, kl


input_dim = 784
hidden_dim = 256
latent_dim = 32
batch_size = 128


def get_inputs():
    p = popcorn_pri
    return [torch.rand(p.jitter_int(batch_size), input_dim)]


def get_init_inputs():
    return [input_dim, hidden_dim, latent_dim]

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

cpp_source = r"""
#include <torch/extension.h>

// Declaration of the function defined in CUDA source
torch::Tensor vae_elbo_forward(
    torch::Tensor x,
    torch::Tensor encoder_weight, torch::Tensor encoder_bias,
    torch::Tensor fc_mu_weight, torch::Tensor fc_mu_bias,
    torch::Tensor fc_logvar_weight, torch::Tensor fc_logvar_bias,
    torch::Tensor decoder1_weight, torch::Tensor decoder1_bias,
    torch::Tensor decoder2_weight, torch::Tensor decoder2_bias,
    torch::Tensor eps,
    torch::Tensor elbo_out,
    torch::Tensor recon_out,
    torch::Tensor kl_out,
    torch::Tensor h_out,
    torch::Tensor mu_out,
    torch::Tensor logvar_out,
    torch::Tensor z_out,
    torch::Tensor h_dec_out,
    torch::Tensor x_recon_out
);
"""

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

// Device functions for activations
__device__ __forceinline__ float relu(float x) {
    return x > 0.0f ? x : 0.0f;
}

__device__ __forceinline__ float sigmoid(float x) {
    return 1.0f / (1.0f + expf(-x));
}

// Fused encoder kernel: Linear + ReLU
__global__ void encoder_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int batch_size, int in_features, int out_features
) {
    int b = blockIdx.x;
    int o = threadIdx.x;
    
    if (b >= batch_size || o >= out_features) return;
    
    float sum = bias[o];
    for (int i = 0; i < in_features; i++) {
        sum += input[b * in_features + i] * weight[o * in_features + i];
    }
    output[b * out_features + o] = relu(sum);
}

// Linear kernel (no activation)
__global__ void linear_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int batch_size, int in_features, int out_features
) {
    int b = blockIdx.x;
    int o = threadIdx.x;
    
    if (b >= batch_size || o >= out_features) return;
    
    float sum = bias[o];
    for (int i = 0; i < in_features; i++) {
        sum += input[b * in_features + i] * weight[o * in_features + i];
    }
    output[b * out_features + o] = sum;
}

// Decoder intermediate: Linear + ReLU
__global__ void decoder_intermediate_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int batch_size, int in_features, int out_features
) {
    int b = blockIdx.x;
    int o = threadIdx.x;
    
    if (b >= batch_size || o >= out_features) return;
    
    float sum = bias[o];
    for (int i = 0; i < in_features; i++) {
        sum += input[b * in_features + i] * weight[o * in_features + i];
    }
    output[b * out_features + o] = relu(sum);
}

// Decoder output: Linear + Sigmoid
__global__ void decoder_output_kernel(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    const float* __restrict__ bias,
    float* __restrict__ output,
    int batch_size, int in_features, int out_features
) {
    int b = blockIdx.x;
    int o = threadIdx.x;
    
    if (b >= batch_size || o >= out_features) return;
    
    float sum = bias[o];
    for (int i = 0; i < in_features; i++) {
        sum += input[b * in_features + i] * weight[o * in_features + i];
    }
    output[b * out_features + o] = sigmoid(sum);
}

// Reparameterization: z = mu + std * eps
__global__ void reparameterize_kernel(
    const float* __restrict__ mu,
    const float* __restrict__ logvar,
    const float* __restrict__ eps,
    float* __restrict__ z,
    int batch_size, int latent_dim
) {
    int b = blockIdx.x;
    int d = threadIdx.x;
    
    if (b >= batch_size || d >= latent_dim) return;
    
    float std_val = expf(0.5f * logvar[b * latent_dim + d]);
    z[b * latent_dim + d] = mu[b * latent_dim + d] + std_val * eps[b * latent_dim + d];
}

// Compute losses: BCE and KL divergence
__global__ void compute_losses_kernel(
    const float* __restrict__ x_recon,
    const float* __restrict__ x,
    const float* __restrict__ mu,
    const float* __restrict__ logvar,
    float* __restrict__ elbo,
    float* __restrict__ recon_loss,
    float* __restrict__ kl_div,
    int batch_size, int input_dim, int latent_dim
) {
    int b = blockIdx.x;
    
    if (b >= batch_size) return;
    
    // Compute reconstruction loss (BCE sum)
    float recon = 0.0f;
    for (int i = 0; i < input_dim; i++) {
        int idx = b * input_dim + i;
        float x_r = x_recon[idx];
        float x_i = x[idx];
        // BCE: -[x*log(x_r) + (1-x)*log(1-x_r)]
        // Clamp for numerical stability
        float x_r_clamped = fmaxf(fminf(x_r, 1.0f - 1e-7f), 1e-7f);
        recon -= x_i * logf(x_r_clamped) + (1.0f - x_i) * logf(1.0f - x_r_clamped);
    }
    
    // Compute KL divergence
    float kl = 0.0f;
    for (int d = 0; d < latent_dim; d++) {
        int idx = b * latent_dim + d;
        float mu_val = mu[idx];
        float logvar_val = logvar[idx];
        // kl = -0.5 * (1 + logvar - mu^2 - exp(logvar))
        kl += -0.5f * (1.0f + logvar_val - mu_val * mu_val - expf(logvar_val));
    }
    
    recon_loss[b] = recon;
    kl_div[b] = kl;
    elbo[b] = -(recon + kl);
}

// Host function
torch::Tensor vae_elbo_forward(
    torch::Tensor x,
    torch::Tensor encoder_weight, torch::Tensor encoder_bias,
    torch::Tensor fc_mu_weight, torch::Tensor fc_mu_bias,
    torch::Tensor fc_logvar_weight, torch::Tensor fc_logvar_bias,
    torch::Tensor decoder1_weight, torch::Tensor decoder1_bias,
    torch::Tensor decoder2_weight, torch::Tensor decoder2_bias,
    torch::Tensor eps,
    torch::Tensor elbo_out,
    torch::Tensor recon_out,
    torch::Tensor kl_out,
    torch::Tensor h_out,
    torch::Tensor mu_out,
    torch::Tensor logvar_out,
    torch::Tensor z_out,
    torch::Tensor h_dec_out,
    torch::Tensor x_recon_out
) {
    int batch_size = x.size(0);
    int input_dim = x.size(1);
    int hidden_dim = encoder_weight.size(0);
    int latent_dim = fc_mu_weight.size(0);
    
    // Encoder: Linear + ReLU
    encoder_kernel<<<batch_size, hidden_dim>>>(
        x.data_ptr<float>(),
        encoder_weight.data_ptr<float>(),
        encoder_bias.data_ptr<float>(),
        h_out.data_ptr<float>(),
        batch_size, input_dim, hidden_dim
    );
    
    // fc_mu: Linear (no activation)
    linear_kernel<<<batch_size, latent_dim>>>(
        h_out.data_ptr<float>(),
        fc_mu_weight.data_ptr<float>(),
        fc_mu_bias.data_ptr<float>(),
        mu_out.data_ptr<float>(),
        batch_size, hidden_dim, latent_dim
    );
    
    // fc_logvar: Linear (no activation)
    linear_kernel<<<batch_size, latent_dim>>>(
        h_out.data_ptr<float>(),
        fc_logvar_weight.data_ptr<float>(),
        fc_logvar_bias.data_ptr<float>(),
        logvar_out.data_ptr<float>(),
        batch_size, hidden_dim, latent_dim
    );
    
    // Reparameterization
    reparameterize_kernel<<<batch_size, latent_dim>>>(
        mu_out.data_ptr<float>(),
        logvar_out.data_ptr<float>(),
        eps.data_ptr<float>(),
        z_out.data_ptr<float>(),
        batch_size, latent_dim
    );
    
    // Decoder: Linear + ReLU
    decoder_intermediate_kernel<<<batch_size, hidden_dim>>>(
        z_out.data_ptr<float>(),
        decoder1_weight.data_ptr<float>(),
        decoder1_bias.data_ptr<float>(),
        h_dec_out.data_ptr<float>(),
        batch_size, latent_dim, hidden_dim
    );
    
    // Decoder output: Linear + Sigmoid
    decoder_output_kernel<<<batch_size, input_dim>>>(
        h_dec_out.data_ptr<float>(),
        decoder2_weight.data_ptr<float>(),
        decoder2_bias.data_ptr<float>(),
        x_recon_out.data_ptr<float>(),
        batch_size, hidden_dim, input_dim
    );
    
    // Compute losses
    compute_losses_kernel<<<batch_size, 1>>>(
        x_recon_out.data_ptr<float>(),
        x.data_ptr<float>(),
        mu_out.data_ptr<float>(),
        logvar_out.data_ptr<float>(),
        elbo_out.data_ptr<float>(),
        recon_out.data_ptr<float>(),
        kl_out.data_ptr<float>(),
        batch_size, input_dim, latent_dim
    );
    
    return elbo_out;
}
"""

# Compile the CUDA extension
vae_module = load_inline(
    name='vae_elbo_cuda_v5',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['vae_elbo_forward'],
    with_cuda=True,
    extra_cuda_cflags=['-O3']
)


class ModelNew(nn.Module):
    def __init__(self, input_dim, hidden_dim, latent_dim):
        super().__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.latent_dim = latent_dim
        
        # Use nn.Linear to match PyTorch initialization exactly
        self.encoder_linear = nn.Linear(input_dim, hidden_dim)
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
        self.decoder1 = nn.Linear(latent_dim, hidden_dim)
        self.decoder2 = nn.Linear(hidden_dim, input_dim)
    
    def forward(self, x: torch.Tensor) -> tuple:
        batch_size = x.size(0)
        
        # Generate random noise for reparameterization
        eps = torch.randn(batch_size, self.latent_dim, device=x.device, dtype=x.dtype)
        
        # Allocate output tensors
        elbo = torch.empty(batch_size, device=x.device, dtype=x.dtype)
        recon = torch.empty(batch_size, device=x.device, dtype=x.dtype)
        kl = torch.empty(batch_size, device=x.device, dtype=x.dtype)
        h = torch.empty(batch_size, self.hidden_dim, device=x.device, dtype=x.dtype)
        mu = torch.empty(batch_size, self.latent_dim, device=x.device, dtype=x.dtype)
        logvar = torch.empty(batch_size, self.latent_dim, device=x.device, dtype=x.dtype)
        z = torch.empty(batch_size, self.latent_dim, device=x.device, dtype=x.dtype)
        h_dec = torch.empty(batch_size, self.hidden_dim, device=x.device, dtype=x.dtype)
        x_recon = torch.empty(batch_size, self.input_dim, device=x.device, dtype=x.dtype)
        
        # Call CUDA kernel with transposed weights (nn.Linear stores weights as [out, in])
        vae_module.vae_elbo_forward(
            x,
            self.encoder_linear.weight, self.encoder_linear.bias,
            self.fc_mu.weight, self.fc_mu.bias,
            self.fc_logvar.weight, self.fc_logvar.bias,
            self.decoder1.weight, self.decoder1.bias,
            self.decoder2.weight, self.decoder2.bias,
            eps,
            elbo, recon, kl,
            h, mu, logvar, z, h_dec, x_recon
        )
        
        return elbo, recon, kl