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 = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void linear_relu_kernel(const float* x, const float* w, const float* b, float* h, int batch, int in_dim, int out_dim) {
    int row = blockIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (row < batch && col < out_dim) {
        float sum = b[col];
        for (int i = 0; i < in_dim; i++) {
            sum += x[row * in_dim + i] * w[col * in_dim + i];
        }
        h[row * out_dim + col] = fmaxf(sum, 0.0f);
    }
}

__global__ void linear_kernel(const float* x, const float* w, const float* b, float* out, int batch, int in_dim, int out_dim) {
    int row = blockIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (row < batch && col < out_dim) {
        float sum = b[col];
        for (int i = 0; i < in_dim; i++) {
            sum += x[row * in_dim + i] * w[col * in_dim + i];
        }
        out[row * out_dim + col] = sum;
    }
}

__global__ void sample_kl_kernel(const float* mu, const float* logvar, const float* eps, float* z, float* kl, int batch, int latent) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < batch * latent) {
        int b = idx / latent;
        float m = mu[idx];
        float lv = logvar[idx];
        float std = expf(0.5f * lv);
        float e = eps[idx];
        z[idx] = m + std * e;
        atomicAdd(&kl[b], -0.5f * (1.0f + lv - m*m - expf(lv)));
    }
}

__global__ void sigmoid_bce_sum(const float* logits, const float* target, float* recon, int batch, int dim) {
    extern __shared__ float sdata[];
    int b = blockIdx.x;
    float sum = 0.0f;
    for (int i = threadIdx.x; i < dim; i += blockDim.x) {
        float lg = logits[b * dim + i];
        float xr = 1.0f / (1.0f + expf(-lg));
        float t = target[b * dim + i];
        sum += -t * logf(xr + 1e-12f) - (1-t) * logf(1.0f - xr + 1e-12f);
    }
    sdata[threadIdx.x] = sum;
    __syncthreads();
    for (int s = blockDim.x/2; s > 0; s >>= 1) {
        if (threadIdx.x < s) sdata[threadIdx.x] += sdata[threadIdx.x + s];
        __syncthreads();
    }
    if (threadIdx.x == 0) recon[b] = sdata[0];
}

std::vector<torch::Tensor> vae_forward(torch::Tensor x, torch::Tensor w_enc, torch::Tensor b_enc,
                          torch::Tensor w_mu, torch::Tensor b_mu, torch::Tensor w_logvar, torch::Tensor b_logvar,
                          torch::Tensor w_dec1, torch::Tensor b_dec1, torch::Tensor w_dec2, torch::Tensor b_dec2, torch::Tensor eps) {
    int batch = x.size(0);
    int input_dim = x.size(1);
    int hidden = w_enc.size(0);
    int latent = w_mu.size(0);
    
    auto options = torch::TensorOptions().device(x.device()).dtype(torch::kFloat32);
    auto h = torch::empty({batch, hidden}, options);
    auto mu = torch::empty({batch, latent}, options);
    auto logvar = torch::empty({batch, latent}, options);
    auto z = torch::empty({batch, latent}, options);
    auto kl = torch::zeros({batch}, options);
    auto h2 = torch::empty({batch, hidden}, options);
    auto logits = torch::empty({batch, input_dim}, options);
    auto recon = torch::empty({batch}, options);
    
    dim3 block(256);
    dim3 grid_enc((hidden + 255)/256, batch);
    linear_relu_kernel<<<grid_enc, block>>>(x.data_ptr<float>(), w_enc.data_ptr<float>(), b_enc.data_ptr<float>(), h.data_ptr<float>(), batch, input_dim, hidden);
    
    dim3 grid_lat((latent + 255)/256, batch);
    linear_kernel<<<grid_lat, block>>>(h.data_ptr<float>(), w_mu.data_ptr<float>(), b_mu.data_ptr<float>(), mu.data_ptr<float>(), batch, hidden, latent);
    linear_kernel<<<grid_lat, block>>>(h.data_ptr<float>(), w_logvar.data_ptr<float>(), b_logvar.data_ptr<float>(), logvar.data_ptr<float>(), batch, hidden, latent);
    
    sample_kl_kernel<<<(batch * latent + 255)/256, 256>>>(mu.data_ptr<float>(), logvar.data_ptr<float>(), eps.data_ptr<float>(), z.data_ptr<float>(), kl.data_ptr<float>(), batch, latent);
    
    dim3 grid_dec((hidden + 255)/256, batch);
    linear_relu_kernel<<<grid_dec, block>>>(z.data_ptr<float>(), w_dec1.data_ptr<float>(), b_dec1.data_ptr<float>(), h2.data_ptr<float>(), batch, latent, hidden);
    
    dim3 grid_out((input_dim + 255)/256, batch);
    linear_kernel<<<grid_out, block>>>(h2.data_ptr<float>(), w_dec2.data_ptr<float>(), b_dec2.data_ptr<float>(), logits.data_ptr<float>(), batch, hidden, input_dim);
    
    sigmoid_bce_sum<<<batch, 256, 256*sizeof(float)>>>(logits.data_ptr<float>(), x.data_ptr<float>(), recon.data_ptr<float>(), batch, input_dim);
    
    auto elbo = -(recon + kl);
    return {elbo, recon, kl};
}
"""

cpp_source = """
std::vector<torch::Tensor> vae_forward(torch::Tensor x, torch::Tensor w_enc, torch::Tensor b_enc,
                          torch::Tensor w_mu, torch::Tensor b_mu, torch::Tensor w_logvar, torch::Tensor b_logvar,
                          torch::Tensor w_dec1, torch::Tensor b_dec1, torch::Tensor w_dec2, torch::Tensor b_dec2, torch::Tensor eps);
"""

module = load_inline(
    name="vae_cuda",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["vae_forward"],
    verbose=True,
    extra_cuda_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):
        w_enc = self.encoder[0].weight
        b_enc = self.encoder[0].bias
        w_mu = self.fc_mu.weight
        b_mu = self.fc_mu.bias
        w_logvar = self.fc_logvar.weight
        b_logvar = self.fc_logvar.bias
        w_dec1 = self.decoder[0].weight
        b_dec1 = self.decoder[0].bias
        w_dec2 = self.decoder[2].weight
        b_dec2 = self.decoder[2].bias
        std = torch.exp(torch.zeros_like(self.fc_mu.bias)) # dummy
        eps = torch.randn_like(std.expand(x.size(0), -1))
        outs = module.vae_forward(x.contiguous(), w_enc, b_enc, w_mu, b_mu, w_logvar, b_logvar, w_dec1, b_dec1, w_dec2, b_dec2, eps)
        return tuple(outs)