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> mu_logvar_forward(torch::Tensor h, torch::Tensor wmu, torch::Tensor bmu, torch::Tensor wlv, torch::Tensor blv);
std::vector<torch::Tensor> zkl_hidden_forward(torch::Tensor mu, torch::Tensor logvar, torch::Tensor eps, torch::Tensor w, torch::Tensor bias);
std::vector<torch::Tensor> recon_elbo_forward(torch::Tensor logits, torch::Tensor x, torch::Tensor kl);
'''

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

__global__ void mu_logvar_kernel(const float* __restrict__ h,
                                 const float* __restrict__ wmu, const float* __restrict__ bmu,
                                 const float* __restrict__ wlv, const float* __restrict__ blv,
                                 float* __restrict__ mu, float* __restrict__ lv) {
    int b = blockIdx.x;
    int t = threadIdx.x;
    int j = t & 31;
    bool is_lv = (t >= 32);
    const float* w = is_lv ? wlv : wmu;
    const float* bias = is_lv ? blv : bmu;
    float sum = bias[j];
    const float* hp = h + b * 256;
    const float* wp = w + j * 256;
    #pragma unroll 4
    for (int k = 0; k < 256; ++k) sum = fmaf(hp[k], wp[k], sum);
    if (is_lv) lv[b * 32 + j] = sum; else mu[b * 32 + j] = sum;
}

__global__ void zkl_hidden_kernel(const float* __restrict__ mu, const float* __restrict__ logvar,
                                  const float* __restrict__ eps, const float* __restrict__ w,
                                  const float* __restrict__ bias, float* __restrict__ dh,
                                  float* __restrict__ kl) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    __shared__ float zsh[32];
    __shared__ float ksh[32];
    if (tid < 32) {
        int idx = b * 32 + tid;
        float m = mu[idx];
        float l = logvar[idx];
        float el = expf(l);
        ksh[tid] = 0.5f * (m * m + el - 1.0f - l);
        zsh[tid] = m + expf(0.5f * l) * eps[idx];
    }
    __syncthreads();
    if (tid < 16) ksh[tid] += ksh[tid + 16];
    __syncthreads();
    if (tid < 8) ksh[tid] += ksh[tid + 8];
    __syncthreads();
    if (tid < 4) ksh[tid] += ksh[tid + 4];
    __syncthreads();
    if (tid < 2) ksh[tid] += ksh[tid + 2];
    __syncthreads();
    if (tid == 0) kl[b] = ksh[0] + ksh[1];

    int j = tid;
    if (j < 256) {
        float sum = bias[j];
        const float* wp = w + j * 32;
        #pragma unroll
        for (int k = 0; k < 32; ++k) sum = fmaf(zsh[k], wp[k], sum);
        dh[b * 256 + j] = sum > 0.0f ? sum : 0.0f;
    }
}

__global__ void recon_elbo_kernel(const float* __restrict__ logits, const float* __restrict__ x,
                                  const float* __restrict__ kl, float* __restrict__ recon,
                                  float* __restrict__ elbo) {
    int b = blockIdx.x;
    int tid = threadIdx.x;
    float sum = 0.0f;
    const int base = b * 784;
    for (int j = tid; j < 784; j += 256) {
        float y = logits[base + j];
        float xv = x[base + j];
        sum += fmaxf(y, 0.0f) - xv * y + log1pf(expf(-fabsf(y)));
    }
    __shared__ float smem[256];
    smem[tid] = sum;
    __syncthreads();
    for (int stride = 128; stride > 0; stride >>= 1) {
        if (tid < stride) smem[tid] += smem[tid + stride];
        __syncthreads();
    }
    if (tid == 0) {
        float r = smem[0];
        recon[b] = r;
        elbo[b] = -(r + kl[b]);
    }
}

std::vector<torch::Tensor> mu_logvar_forward(torch::Tensor h, torch::Tensor wmu, torch::Tensor bmu, torch::Tensor wlv, torch::Tensor blv) {
    int B = h.size(0);
    auto mu = torch::empty({B, 32}, h.options());
    auto lv = torch::empty({B, 32}, h.options());
    mu_logvar_kernel<<<B, 64>>>(h.data_ptr<float>(), wmu.data_ptr<float>(), bmu.data_ptr<float>(),
                                wlv.data_ptr<float>(), blv.data_ptr<float>(), mu.data_ptr<float>(), lv.data_ptr<float>());
    return {mu, lv};
}

std::vector<torch::Tensor> zkl_hidden_forward(torch::Tensor mu, torch::Tensor logvar, torch::Tensor eps, torch::Tensor w, torch::Tensor bias) {
    int B = mu.size(0);
    auto dh = torch::empty({B, 256}, mu.options());
    auto kl = torch::empty({B}, mu.options());
    zkl_hidden_kernel<<<B, 256>>>(mu.data_ptr<float>(), logvar.data_ptr<float>(), eps.data_ptr<float>(),
                                  w.data_ptr<float>(), bias.data_ptr<float>(), dh.data_ptr<float>(), kl.data_ptr<float>());
    return {dh, kl};
}

std::vector<torch::Tensor> recon_elbo_forward(torch::Tensor logits, torch::Tensor x, torch::Tensor kl) {
    int B = logits.size(0);
    auto recon = torch::empty({B}, logits.options());
    auto elbo = torch::empty({B}, logits.options());
    recon_elbo_kernel<<<B, 256>>>(logits.data_ptr<float>(), x.data_ptr<float>(), kl.data_ptr<float>(), recon.data_ptr<float>(), elbo.data_ptr<float>());
    return {elbo, recon};
}
'''

_ext = load_inline(name='vae_elbo_ext_final2', cpp_sources=cpp_src, cuda_sources=cuda_src,
                   functions=['mu_logvar_forward','zkl_hidden_forward','recon_elbo_forward'], verbose=False,
                   extra_cuda_cflags=['-O3', '--use_fast_math'], extra_cflags=['-O3'])

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:
        h = F.linear(x, self.encoder[0].weight, self.encoder[0].bias)
        h = torch.relu(h)
        mu, logvar = _ext.mu_logvar_forward(h, self.fc_mu.weight, self.fc_mu.bias, self.fc_logvar.weight, self.fc_logvar.bias)
        eps = torch.randn_like(logvar)
        dh, kl = _ext.zkl_hidden_forward(mu, logvar, eps, self.decoder[0].weight, self.decoder[0].bias)
        logits = F.linear(dh, self.decoder[2].weight, self.decoder[2].bias)
        elbo, recon = _ext.recon_elbo_forward(logits, x, kl)
        return elbo, recon, kl