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_src = r"""
#include <torch/extension.h>
#include <vector>

std::vector<torch::Tensor> reparam_kl(torch::Tensor mu, torch::Tensor logvar, torch::Tensor eps);
std::vector<torch::Tensor> bce_elbo(torch::Tensor logits, torch::Tensor x, torch::Tensor kl);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <vector>
#include <math.h>

#ifndef CHECK_CUDA
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
#endif

__global__ void reparam_kl_kernel(const float* __restrict__ mu,
                                  const float* __restrict__ logvar,
                                  const float* __restrict__ eps,
                                  float* __restrict__ z,
                                  float* __restrict__ kl,
                                  int L) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    extern __shared__ float smem[];
    float acc = 0.0f;
    int base = b * L;

    for (int j = tid; j < L; j += blockDim.x) {
        float m = mu[base + j];
        float lv = logvar[base + j];
        float ev = expf(lv);
        float stdv = expf(0.5f * lv);
        z[base + j] = m + stdv * eps[base + j];
        acc += -0.5f * (1.0f + lv - m * m - ev);
    }

    smem[tid] = acc;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }
    if (tid == 0) kl[b] = smem[0];
}

__global__ void bce_elbo_kernel(const float* __restrict__ logits,
                                const float* __restrict__ x,
                                const float* __restrict__ kl,
                                float* __restrict__ elbo,
                                float* __restrict__ recon,
                                int D) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    extern __shared__ float smem[];
    float acc = 0.0f;
    int base = b * D;

    for (int j = tid; j < D; j += blockDim.x) {
        float v = logits[base + j];
        float t = x[base + j];
        // Stable binary cross entropy with logits.  This is numerically
        // equivalent to BCE(sigmoid(v), t) for the value ranges produced here,
        // while avoiding materializing the sigmoid/probability tensor.
        float av = fabsf(v);
        float mx = v > 0.0f ? v : 0.0f;
        acc += mx - v * t + log1pf(expf(-av));
    }

    smem[tid] = acc;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }
    if (tid == 0) {
        float r = smem[0];
        recon[b] = r;
        elbo[b] = -(r + kl[b]);
    }
}

std::vector<torch::Tensor> reparam_kl(torch::Tensor mu, torch::Tensor logvar, torch::Tensor eps) {
    CHECK_CUDA(mu); CHECK_CUDA(logvar); CHECK_CUDA(eps);
    c10::cuda::CUDAGuard device_guard(mu.device());
    const int B = (int)mu.size(0);
    const int L = (int)mu.size(1);
    auto z = torch::empty_like(mu);
    auto kl = torch::empty({B}, mu.options());
    const int threads = 64;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    reparam_kl_kernel<<<B, threads, threads * sizeof(float), stream>>>(
        mu.data_ptr<float>(), logvar.data_ptr<float>(), eps.data_ptr<float>(),
        z.data_ptr<float>(), kl.data_ptr<float>(), L);
    return {z, kl};
}

std::vector<torch::Tensor> bce_elbo(torch::Tensor logits, torch::Tensor x, torch::Tensor kl) {
    CHECK_CUDA(logits); CHECK_CUDA(x); CHECK_CUDA(kl);
    c10::cuda::CUDAGuard device_guard(logits.device());
    const int B = (int)logits.size(0);
    const int D = (int)logits.size(1);
    auto recon = torch::empty({B}, logits.options());
    auto elbo = torch::empty({B}, logits.options());
    const int threads = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    bce_elbo_kernel<<<B, threads, threads * sizeof(float), stream>>>(
        logits.data_ptr<float>(), x.data_ptr<float>(), kl.data_ptr<float>(),
        elbo.data_ptr<float>(), recon.data_ptr<float>(), D);
    return {elbo, recon};
}
"""

_vae_elbo_ext = load_inline(
    name="vae_elbo_fused_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["reparam_kl", "bce_elbo"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)


class ModelNew(nn.Module):
    """CUDA-fused forward for the VAE ELBO model."""

    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:
        # Keep the same random draw location/shape as the reference, but fuse
        # reparameterization + KL and final BCE reduction + ELBO construction.
        h = F.linear(x, self.encoder[0].weight, self.encoder[0].bias)
        h = torch.relu(h)

        mu = F.linear(h, self.fc_mu.weight, self.fc_mu.bias)
        logvar = F.linear(h, self.fc_logvar.weight, self.fc_logvar.bias)

        eps = torch.randn_like(logvar)
        z, kl = _vae_elbo_ext.reparam_kl(mu, logvar, eps)

        dh = F.linear(z, self.decoder[0].weight, self.decoder[0].bias)
        dh = torch.relu(dh)
        logits = F.linear(dh, self.decoder[2].weight, self.decoder[2].bias)

        elbo, recon = _vae_elbo_ext.bce_elbo(logits, x, kl)
        return elbo, recon, kl