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
from torch.utils.cpp_extension import load_inline

cuda_source = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

__global__ void elbo_kernel(const float* x, const float* x_recon, const float* mu, const float* logvar, float* elbo, float* recon, float* kl, int batch, int input_dim, int latent_dim) {
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch) return;
    float r = 0.0f;
    for (int i = 0; i < input_dim; i++) {
        float xr = x_recon[b * input_dim + i];
        float xx = x[b * input_dim + i];
        r += - (xx * logf(xr + 1e-12f) + (1 - xx) * logf(1 - xr + 1e-12f));
    }
    float k = 0.0f;
    for (int i = 0; i < latent_dim; i++) {
        float m = mu[b * latent_dim + i];
        float lv = logvar[b * latent_dim + i];
        k += -0.5f * (1.0f + lv - m*m - expf(lv));
    }
    recon[b] = r;
    kl[b] = k;
    elbo[b] = -(r + k);
}

std::vector<torch::Tensor> compute_elbo(torch::Tensor x, torch::Tensor x_recon, torch::Tensor mu, torch::Tensor logvar) {
    int batch = x.size(0);
    int input_dim = x.size(1);
    int latent_dim = mu.size(1);
    auto elbo = torch::empty({batch}, x.options());
    auto recon = torch::empty({batch}, x.options());
    auto kl = torch::empty({batch}, x.options());
    elbo_kernel<<<(batch + 255)/256, 256>>>(x.data_ptr<float>(), x_recon.data_ptr<float>(), mu.data_ptr<float>(), logvar.data_ptr<float>(), elbo.data_ptr<float>(), recon.data_ptr<float>(), kl.data_ptr<float>(), batch, input_dim, latent_dim);
    return {elbo, recon, kl};
}
'''

cpp_source = r'''
std::vector<torch::Tensor> compute_elbo(torch::Tensor x, torch::Tensor x_recon, torch::Tensor mu, torch::Tensor logvar);
'''

module = load_inline(
    name='elbo_fuse',
    cpp_sources=[cpp_source],
    cuda_sources=[cuda_source],
    functions=['compute_elbo'],
    verbose=False
)

class ModelNew(nn.Module):
    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):
        h = self.encoder(x)
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        z = mu + std * eps
        x_recon = self.decoder(z)
        outs = module.compute_elbo(x, x_recon, mu, logvar)
        return tuple(outs)