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_src = r'''
#include <torch/extension.h>
#include <vector>
std::vector<torch::Tensor> cvae_forward(torch::Tensor x, torch::Tensor y, torch::Tensor enc_wt, torch::Tensor enc_b, torch::Tensor mu_wt, torch::Tensor mu_b, torch::Tensor logvar_wt, torch::Tensor logvar_b, torch::Tensor dec0_wt, torch::Tensor dec0_b, torch::Tensor dec2_wt, torch::Tensor dec2_b);
'''

_cuda_src = 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__ __launch_bounds__(256, 2) void cvae_main_kernel(
    const float* __restrict__ x,
    const float* __restrict__ y,
    const float* __restrict__ eps,
    const float* __restrict__ enc_wt,
    const float* __restrict__ enc_b,
    const float* __restrict__ mu_wt,
    const float* __restrict__ mu_b,
    const float* __restrict__ logvar_wt,
    const float* __restrict__ logvar_b,
    const float* __restrict__ dec0_wt,
    const float* __restrict__ dec0_b,
    const float* __restrict__ dec2_wt,
    const float* __restrict__ dec2_b,
    float* __restrict__ partial,
    int B) {

    const int b = blockIdx.x;
    const int tid = threadIdx.x;

    __shared__ float h[HIDDEN_DIM];
    __shared__ float z[LATENT_DIM];
    __shared__ float dh[HIDDEN_DIM];
    __shared__ float red[256];

    const float* xb = x + b * INPUT_DIM;
    const float* yb = y + b * LABEL_DIM;

    if (tid < HIDDEN_DIM) {
        float s = enc_b[tid];
        #pragma unroll 4
        for (int i = 0; i < INPUT_DIM; ++i) {
            s = fmaf(xb[i], enc_wt[i * HIDDEN_DIM + tid], s);
        }
        #pragma unroll
        for (int l = 0; l < LABEL_DIM; ++l) {
            s = fmaf(yb[l], enc_wt[(INPUT_DIM + l) * HIDDEN_DIM + tid], s);
        }
        h[tid] = s > 0.0f ? s : 0.0f;
    }
    __syncthreads();

    float kval = 0.0f;
    if (tid < LATENT_DIM) {
        float sm = mu_b[tid];
        float sl = logvar_b[tid];
        #pragma unroll 4
        for (int j = 0; j < HIDDEN_DIM; ++j) {
            float hv = h[j];
            sm = fmaf(hv, mu_wt[j * LATENT_DIM + tid], sm);
            sl = fmaf(hv, logvar_wt[j * LATENT_DIM + tid], sl);
        }
        z[tid] = sm + expf(0.5f * sl) * eps[b * LATENT_DIM + tid];
        kval = -0.5f * (1.0f + sl - sm * sm - expf(sl));
    }

    red[tid] = (tid < LATENT_DIM) ? kval : 0.0f;
    __syncthreads();
    for (int off = 128; off > 0; off >>= 1) {
        if (tid < off) red[tid] += red[tid + off];
        __syncthreads();
    }
    if (tid == 0) partial[B + b] = red[0];
    __syncthreads();

    if (tid < HIDDEN_DIM) {
        float s = dec0_b[tid];
        #pragma unroll
        for (int k = 0; k < LATENT_DIM; ++k) {
            s = fmaf(z[k], dec0_wt[k * HIDDEN_DIM + tid], s);
        }
        #pragma unroll
        for (int l = 0; l < LABEL_DIM; ++l) {
            s = fmaf(yb[l], dec0_wt[(LATENT_DIM + l) * HIDDEN_DIM + tid], s);
        }
        dh[tid] = s > 0.0f ? s : 0.0f;
    }
    __syncthreads();

    float r = 0.0f;
    if (tid < INPUT_DIM) {
        float s = dec2_b[tid];
        #pragma unroll 4
        for (int j = 0; j < HIDDEN_DIM; ++j) {
            s = fmaf(dh[j], dec2_wt[j * INPUT_DIM + tid], s);
        }
        float d = s - xb[tid];
        r = d * d;
    }

    red[tid] = r;
    __syncthreads();
    for (int off = 128; off > 0; off >>= 1) {
        if (tid < off) red[tid] += red[tid + off];
        __syncthreads();
    }
    if (tid == 0) partial[b] = red[0];
}

__global__ void cvae_final_reduce_kernel(const float* __restrict__ partial, float* __restrict__ out, int B) {
    int tid = threadIdx.x;
    __shared__ float red[256];

    float s = 0.0f;
    for (int i = tid; i < B; i += blockDim.x) s += partial[i];
    red[tid] = s;
    __syncthreads();
    for (int off = 128; off > 0; off >>= 1) {
        if (tid < off) red[tid] += red[tid + off];
        __syncthreads();
    }
    float recon = red[0] / (float)B;

    s = 0.0f;
    for (int i = tid; i < B; i += blockDim.x) s += partial[B + i];
    red[tid] = s;
    __syncthreads();
    for (int off = 128; off > 0; off >>= 1) {
        if (tid < off) red[tid] += red[tid + off];
        __syncthreads();
    }

    if (tid == 0) {
        float kl = red[0] / (float)B;
        out[0] = recon + kl;
        out[1] = recon;
        out[2] = kl;
    }
}

std::vector<torch::Tensor> cvae_forward(
    torch::Tensor x,
    torch::Tensor y,
    torch::Tensor enc_wt,
    torch::Tensor enc_b,
    torch::Tensor mu_wt,
    torch::Tensor mu_b,
    torch::Tensor logvar_wt,
    torch::Tensor logvar_b,
    torch::Tensor dec0_wt,
    torch::Tensor dec0_b,
    torch::Tensor dec2_wt,
    torch::Tensor dec2_b) {

    int B = (int)x.size(0);
    auto opts = x.options();
    auto eps = torch::randn({B, LATENT_DIM}, opts);
    auto partial = torch::empty({2, B}, opts);
    auto out = torch::empty({3}, opts);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    cvae_main_kernel<<<B, 256, 0, stream>>>(
        x.data_ptr<float>(), y.data_ptr<float>(), eps.data_ptr<float>(),
        enc_wt.data_ptr<float>(), enc_b.data_ptr<float>(),
        mu_wt.data_ptr<float>(), mu_b.data_ptr<float>(),
        logvar_wt.data_ptr<float>(), logvar_b.data_ptr<float>(),
        dec0_wt.data_ptr<float>(), dec0_b.data_ptr<float>(),
        dec2_wt.data_ptr<float>(), dec2_b.data_ptr<float>(),
        partial.data_ptr<float>(), B);

    cvae_final_reduce_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_final',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['cvae_forward'],
    verbose=False,
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    extra_cflags=['-O3'],
)

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),
        )
        self.register_buffer('_enc_wt', torch.empty(input_dim + label_dim, hidden_dim), persistent=False)
        self.register_buffer('_mu_wt', torch.empty(hidden_dim, latent_dim), persistent=False)
        self.register_buffer('_logvar_wt', torch.empty(hidden_dim, latent_dim), persistent=False)
        self.register_buffer('_dec0_wt', torch.empty(latent_dim + label_dim, hidden_dim), persistent=False)
        self.register_buffer('_dec2_wt', torch.empty(hidden_dim, input_dim), persistent=False)
        self._refresh_cache()

    def _refresh_cache(self):
        with torch.no_grad():
            self._enc_wt = self.enc[0].weight.detach().t().contiguous()
            self._mu_wt = self.enc_mu.weight.detach().t().contiguous()
            self._logvar_wt = self.enc_logvar.weight.detach().t().contiguous()
            self._dec0_wt = self.dec[0].weight.detach().t().contiguous()
            self._dec2_wt = self.dec[2].weight.detach().t().contiguous()

    def load_state_dict(self, state_dict, strict=True, assign=False):
        result = super().load_state_dict(state_dict, strict=strict, assign=assign)
        self._refresh_cache()
        return result

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> tuple:
        return tuple(_ext.cvae_forward(
            x, y,
            self._enc_wt, self.enc[0].bias,
            self._mu_wt, self.enc_mu.bias,
            self._logvar_wt, self.enc_logvar.bias,
            self._dec0_wt, self.dec[0].bias,
            self._dec2_wt, self.dec[2].bias,
        ))