import torch
import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    """
    Conditional Variational Autoencoder (CVAE) with full encoder–decoder
    forward pass and loss computation.  The encoder conditions on both
    input x and label y; the decoder conditions on z and y.  Returns
    reconstruction loss + KL divergence.
    """

    def __init__(self, input_dim, label_dim, hidden_dim, latent_dim):
        super().__init__()
        # Encoder: q(z | x, y)
        self.enc = nn.Sequential(
            nn.Linear(input_dim + label_dim, hidden_dim),
            nn.ReLU(),
        )
        self.enc_mu = nn.Linear(hidden_dim, latent_dim)
        self.enc_logvar = nn.Linear(hidden_dim, latent_dim)

        # Decoder: p(x | z, y)
        self.dec = nn.Sequential(
            nn.Linear(latent_dim + label_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim),
        )

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> tuple:
        """
        Args:
            x: (B, input_dim)  – input data
            y: (B, label_dim)  – conditioning label (one-hot)
        Returns:
            (loss, recon_loss, kl_loss):
                loss:       scalar – total CVAE loss
                recon_loss: scalar
                kl_loss:    scalar
        """
        h = self.enc(torch.cat([x, y], dim=-1))
        mu = self.enc_mu(h)
        logvar = self.enc_logvar(h)

        std = torch.exp(0.5 * logvar)
        z = mu + std * torch.randn_like(std)

        x_recon = self.dec(torch.cat([z, y], dim=-1))
        recon_loss = F.mse_loss(x_recon, x, reduction="sum") / x.shape[0]
        kl_loss = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()).sum() / x.shape[0]
        loss = recon_loss + kl_loss
        return loss, recon_loss, kl_loss


input_dim = 256
label_dim = 10
hidden_dim = 128
latent_dim = 32
batch_size = 128


def get_inputs():
    p = popcorn_pri
    x = torch.randn(p.jitter_int(batch_size), input_dim)
    y = torch.zeros(p.jitter_int(batch_size), label_dim)
    y.scatter_(1, torch.randint(0, label_dim, (p.jitter_int(batch_size), 1)), 1.0)
    return [x, y]


def get_init_inputs():
    return [input_dim, label_dim, hidden_dim, latent_dim]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

cpp_source = r'''
#include <torch/extension.h>
#include <vector>
std::vector<torch::Tensor> cvae_encoder(torch::Tensor x, torch::Tensor y, torch::Tensor w_enc, torch::Tensor b_enc, torch::Tensor w_mu, torch::Tensor b_mu, torch::Tensor w_lv, torch::Tensor b_lv);
std::vector<torch::Tensor> cvae_decoder_loss(torch::Tensor x, torch::Tensor y, torch::Tensor eps, torch::Tensor mu, torch::Tensor logvar, torch::Tensor w_dec0, torch::Tensor b_dec0, torch::Tensor w_dec2, torch::Tensor b_dec2);
'''

cuda_source = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#define INPUT_DIM 256
#define LABEL_DIM 10
#define HIDDEN_DIM 128
#define LATENT_DIM 32

__global__ void encoder_kernel(const float* __restrict__ x, const float* __restrict__ y,
                               const float* __restrict__ w_enc, const float* __restrict__ b_enc,
                               const float* __restrict__ w_mu, const float* __restrict__ b_mu,
                               const float* __restrict__ w_lv, const float* __restrict__ b_lv,
                               float* __restrict__ mu, float* __restrict__ logvar) {
    int row = blockIdx.x, tid = threadIdx.x;
    __shared__ float h[HIDDEN_DIM];
    float sum = b_enc[tid];
    const float* wx = w_enc + tid * (INPUT_DIM + LABEL_DIM);
    const float* xr = x + row * INPUT_DIM;
    #pragma unroll 4
    for (int k = 0; k < INPUT_DIM; ++k) sum += xr[k] * wx[k];
    const float* yr = y + row * LABEL_DIM;
    #pragma unroll
    for (int k = 0; k < LABEL_DIM; ++k) sum += yr[k] * wx[INPUT_DIM + k];
    h[tid] = fmaxf(sum, 0.0f);
    __syncthreads();
    if (tid < LATENT_DIM) {
        float sm = b_mu[tid], sl = b_lv[tid];
        const float* wm = w_mu + tid * HIDDEN_DIM;
        const float* wl = w_lv + tid * HIDDEN_DIM;
        #pragma unroll 4
        for (int k = 0; k < HIDDEN_DIM; ++k) { float hv = h[k]; sm += hv * wm[k]; sl += hv * wl[k]; }
        mu[row * LATENT_DIM + tid] = sm;
        logvar[row * LATENT_DIM + tid] = sl;
    }
}

__global__ void decoder_loss_kernel(const float* __restrict__ x, const float* __restrict__ y,
                                    const float* __restrict__ eps, const float* __restrict__ mu, const float* __restrict__ logvar,
                                    const float* __restrict__ w_dec0, const float* __restrict__ b_dec0,
                                    const float* __restrict__ w_dec2, const float* __restrict__ b_dec2,
                                    float* __restrict__ partial) {
    int row = blockIdx.x, tid = threadIdx.x;
    __shared__ float z[LATENT_DIM];
    __shared__ float dh[HIDDEN_DIM];
    __shared__ float red_mse[INPUT_DIM];
    __shared__ float red_kl[INPUT_DIM];
    float kl = 0.0f;
    if (tid < LATENT_DIM) {
        float m = mu[row * LATENT_DIM + tid];
        float lv = logvar[row * LATENT_DIM + tid];
        z[tid] = m + expf(0.5f * lv) * eps[row * LATENT_DIM + tid];
        kl = -0.5f * (1.0f + lv - m * m - expf(lv));
    }
    __syncthreads();
    if (tid < HIDDEN_DIM) {
        float sum = b_dec0[tid];
        const float* wd = w_dec0 + tid * (LATENT_DIM + LABEL_DIM);
        #pragma unroll
        for (int k = 0; k < LATENT_DIM; ++k) sum += z[k] * wd[k];
        const float* yr = y + row * LABEL_DIM;
        #pragma unroll
        for (int k = 0; k < LABEL_DIM; ++k) sum += yr[k] * wd[LATENT_DIM + k];
        dh[tid] = fmaxf(sum, 0.0f);
    }
    __syncthreads();
    float mse = 0.0f;
    if (tid < INPUT_DIM) {
        float sum = b_dec2[tid];
        const float* w2 = w_dec2 + tid * HIDDEN_DIM;
        #pragma unroll 4
        for (int k = 0; k < HIDDEN_DIM; ++k) sum += dh[k] * w2[k];
        float d = sum - x[row * INPUT_DIM + tid];
        mse = d * d;
    }
    red_mse[tid] = mse;
    red_kl[tid] = kl;
    __syncthreads();
    for (int offset = 128; offset > 0; offset >>= 1) {
        if (tid < offset) { red_mse[tid] += red_mse[tid + offset]; red_kl[tid] += red_kl[tid + offset]; }
        __syncthreads();
    }
    if (tid == 0) { partial[row] = red_mse[0]; partial[gridDim.x + row] = red_kl[0]; }
}

__global__ void finalize_kernel(const float* __restrict__ partial, float* __restrict__ out, int B) {
    int tid = threadIdx.x;
    __shared__ float sm[256];
    __shared__ float sk[256];
    float m = 0.0f, k = 0.0f;
    if (tid < B) { m = partial[tid]; k = partial[B + tid]; }
    sm[tid] = m; sk[tid] = k;
    __syncthreads();
    for (int off = 128; off > 0; off >>= 1) {
        if (tid < off) { sm[tid] += sm[tid+off]; sk[tid] += sk[tid+off]; }
        __syncthreads();
    }
    if (tid == 0) {
        float invB = 1.0f / (float)B;
        float recon = sm[0] * invB;
        float kl = sk[0] * invB;
        out[0] = recon + kl; out[1] = recon; out[2] = kl;
    }
}

std::vector<torch::Tensor> cvae_encoder(torch::Tensor x, torch::Tensor y, torch::Tensor w_enc, torch::Tensor b_enc, torch::Tensor w_mu, torch::Tensor b_mu, torch::Tensor w_lv, torch::Tensor b_lv) {
    auto mu = torch::empty({x.size(0), LATENT_DIM}, x.options());
    auto logvar = torch::empty({x.size(0), LATENT_DIM}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    encoder_kernel<<<x.size(0), HIDDEN_DIM, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>(), w_enc.data_ptr<float>(), b_enc.data_ptr<float>(), w_mu.data_ptr<float>(), b_mu.data_ptr<float>(), w_lv.data_ptr<float>(), b_lv.data_ptr<float>(), mu.data_ptr<float>(), logvar.data_ptr<float>());
    return {mu, logvar};
}
std::vector<torch::Tensor> cvae_decoder_loss(torch::Tensor x, torch::Tensor y, torch::Tensor eps, torch::Tensor mu, torch::Tensor logvar, torch::Tensor w_dec0, torch::Tensor b_dec0, torch::Tensor w_dec2, torch::Tensor b_dec2) {
    int B = x.size(0);
    auto partial = torch::empty({2 * B}, x.options());
    auto out = torch::empty({3}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    decoder_loss_kernel<<<B, INPUT_DIM, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>(), eps.data_ptr<float>(), mu.data_ptr<float>(), logvar.data_ptr<float>(), w_dec0.data_ptr<float>(), b_dec0.data_ptr<float>(), w_dec2.data_ptr<float>(), b_dec2.data_ptr<float>(), partial.data_ptr<float>());
    finalize_kernel<<<1, 256, 0, stream>>>(partial.data_ptr<float>(), out.data_ptr<float>(), B);
    return {out.select(0,0), out.select(0,1), out.select(0,2)};
}
'''

_ext = load_inline(name='cvae_fused_ext_v2', cpp_sources=cpp_source, cuda_sources=cuda_source,
                   functions=['cvae_encoder', 'cvae_decoder_loss'], extra_cuda_cflags=['-O3','--use_fast_math'], extra_cflags=['-O3'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, input_dim, label_dim, hidden_dim, latent_dim):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(input_dim + label_dim, hidden_dim), nn.ReLU())
        self.enc_mu = nn.Linear(hidden_dim, latent_dim)
        self.enc_logvar = nn.Linear(hidden_dim, latent_dim)
        self.dec = nn.Sequential(nn.Linear(latent_dim + label_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, input_dim))
    def forward(self, x: torch.Tensor, y: torch.Tensor) -> tuple:
        enc0 = self.enc[0]
        mu, logvar = _ext.cvae_encoder(x, y, enc0.weight, enc0.bias, self.enc_mu.weight, self.enc_mu.bias, self.enc_logvar.weight, self.enc_logvar.bias)
        eps = torch.randn_like(mu)
        return tuple(_ext.cvae_decoder_loss(x, y, eps, mu, logvar, self.dec[0].weight, self.dec[0].bias, self.dec[2].weight, self.dec[2].bias))