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

cpp_src = r'''
#include <torch/extension.h>
#include <vector>
std::vector<torch::Tensor> vae_forward(torch::Tensor x,
    torch::Tensor ew, torch::Tensor eb,
    torch::Tensor mw, torch::Tensor mb,
    torch::Tensor lw, torch::Tensor lb,
    torch::Tensor d0w, torch::Tensor d0b,
    torch::Tensor d2w, torch::Tensor d2b);
'''

cuda_src = r'''
#include <torch/extension.h>
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <vector>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void sample_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;
    __shared__ float sm[32];
    float sum = 0.0f;
    int idx = b * L + tid;
    if (tid < L) {
        float m = mu[idx];
        float lv = logvar[idx];
        float st = expf(0.5f * lv);
        z[idx] = m + st * eps[idx];
        sum = -0.5f * (1.0f + lv - m*m - expf(lv));
    }
    sm[tid] = sum;
    __syncthreads();
    for (int s = 16; s > 0; s >>= 1) {
        if (tid < s) sm[tid] += sm[tid+s];
        __syncthreads();
    }
    if (tid == 0) kl[b] = sm[0];
}

__device__ __forceinline__ float softplus_stable(float v) {
    float av = fabsf(v);
    return fmaxf(v, 0.0f) + log1pf(expf(-av));
}

__global__ void recon_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 rsum = 0.0f;
    int baseD = b * D;
    for (int j = tid; j < D; j += blockDim.x) {
        float l = logits[baseD + j];
        float xv = x[baseD + j];
        rsum += softplus_stable(l) - xv * l;
    }
    smem[tid] = rsum;
    __syncthreads();
    for (int s = blockDim.x/2; 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> vae_forward(torch::Tensor x,
    torch::Tensor ew, torch::Tensor eb,
    torch::Tensor mw, torch::Tensor mb,
    torch::Tensor lw, torch::Tensor lb,
    torch::Tensor d0w, torch::Tensor d0b,
    torch::Tensor d2w, torch::Tensor d2b) {
    auto h = at::linear(x, ew, eb);
    h.relu_();
    auto mu = at::linear(h, mw, mb);
    auto logvar = at::linear(h, lw, lb);

    int B = x.size(0);
    int D = x.size(1);
    int L = mu.size(1);
    auto opts = x.options();

    auto eps = at::randn_like(logvar);
    auto z = torch::empty_like(mu);
    auto kl = torch::empty({B}, opts);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    sample_kl_kernel<<<B, 32, 0, stream>>>(mu.data_ptr<float>(), logvar.data_ptr<float>(), eps.data_ptr<float>(), z.data_ptr<float>(), kl.data_ptr<float>(), L);

    auto dh = at::linear(z, d0w, d0b);
    dh.relu_();
    auto logits = at::linear(dh, d2w, d2b);

    auto elbo = torch::empty({B}, opts);
    auto recon = torch::empty({B}, opts);
    const int threads = 256;
    recon_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, kl};
}
'''

_ext = load_inline(name='vae_elbo_ext', cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=['vae_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'], 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: torch.Tensor) -> tuple:
        return tuple(_ext.vae_forward(
            x.contiguous(),
            self.encoder[0].weight, self.encoder[0].bias,
            self.fc_mu.weight, self.fc_mu.bias,
            self.fc_logvar.weight, self.fc_logvar.bias,
            self.decoder[0].weight, self.decoder[0].bias,
            self.decoder[2].weight, self.decoder[2].bias,
        ))