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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

cuda_source = r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <curand_kernel.h>

__global__ void cvae_forward_kernel(
    const float* __restrict__ x,
    const float* __restrict__ y,
    const float* __restrict__ enc_weight,
    const float* __restrict__ enc_bias,
    const float* __restrict__ enc_mu_weight,
    const float* __restrict__ enc_mu_bias,
    const float* __restrict__ enc_logvar_weight,
    const float* __restrict__ enc_logvar_bias,
    const float* __restrict__ dec_weight1,
    const float* __restrict__ dec_bias1,
    const float* __restrict__ dec_weight2,
    const float* __restrict__ dec_bias2,
    float* __restrict__ loss_out,
    float* __recon_loss_out,
    float* __kl_loss_out,
    int batch_size,
    int input_dim,
    int label_dim,
    int hidden_dim,
    int latent_dim,
    unsigned long long seed
) {
    int b = blockIdx.x;
    if (b >= batch_size) return;

    extern __shared__ float smem[];

    // Simple implementation: each block handles one sample, using shared mem for intermediates
    float* h = smem;
    float* mu = &smem[hidden_dim];
    float* logvar = &smem[hidden_dim + latent_dim];
    float* z = &smem[hidden_dim + 2*latent_dim];
    float* h2 = &smem[hidden_dim + 3*latent_dim];
    float* x_recon = &smem[hidden_dim + 3*latent_dim + hidden_dim];

    // Encoder first linear + ReLU
    for (int i = threadIdx.x; i < hidden_dim; i += blockDim.x) {
        float sum = enc_bias[i];
        for (int j = 0; j < input_dim + label_dim; j++) {
            float val = (j < input_dim) ? x[b * input_dim + j] : y[b * label_dim + (j - input_dim)];
            sum += val * enc_weight[i * (input_dim + label_dim) + j];
        }
        h[i] = fmaxf(sum, 0.0f);
    }
    __syncthreads();

    // mu and logvar
    for (int i = threadIdx.x; i < latent_dim; i += blockDim.x) {
        float sum_mu = enc_mu_bias[i];
        float sum_logvar = enc_logvar_bias[i];
        for (int j = 0; j < hidden_dim; j++) {
            sum_mu += h[j] * enc_mu_weight[i * hidden_dim + j];
            sum_logvar += h[j] * enc_logvar_weight[i * hidden_dim + j];
        }
        mu[i] = sum_mu;
        logvar[i] = sum_logvar;
    }
    __syncthreads();

    // Sample z
    for (int i = threadIdx.x; i < latent_dim; i += blockDim.x) {
        float std = expf(0.5f * logvar[i]);
        curandState state;
        curand_init(seed + b * latent_dim + i, 0, 0, &state);
        float eps = curand_normal(&state);
        z[i] = mu[i] + std * eps;
    }
    __syncthreads();

    // Decoder
    for (int i = threadIdx.x; i < hidden_dim; i += blockDim.x) {
        float sum = dec_bias1[i];
        for (int j = 0; j < latent_dim + label_dim; j++) {
            float val = (j < latent_dim) ? z[j] : y[b * label_dim + (j - latent_dim)];
            sum += val * dec_weight1[i * (latent_dim + label_dim) + j];
        }
        h2[i] = fmaxf(sum, 0.0f);
    }
    __syncthreads();

    for (int i = threadIdx.x; i < input_dim; i += blockDim.x) {
        float sum = dec_bias2[i];
        for (int j = 0; j < hidden_dim; j++) {
            sum += h2[j] * dec_weight2[i * hidden_dim + j];
        }
        x_recon[i] = sum;
    }
    __syncthreads();

    // Compute per-sample recon and kl (will reduce later)
    __shared__ float partial_recon;
    __shared__ float partial_kl;
    if (threadIdx.x == 0) {
        partial_recon = 0.0f;
        partial_kl = 0.0f;
    }
    __syncthreads();

    float local_recon = 0.0f;
    float local_kl = 0.0f;
    for (int i = threadIdx.x; i < input_dim; i += blockDim.x) {
        float diff = x_recon[i] - x[b * input_dim + i];
        local_recon += diff * diff;
    }
    for (int i = threadIdx.x; i < latent_dim; i += blockDim.x) {
        local_kl += -0.5f * (1.0f + logvar[i] - mu[i]*mu[i] - expf(logvar[i]));
    }
    atomicAdd(&partial_recon, local_recon);
    atomicAdd(&partial_kl, local_kl);
    __syncthreads();

    if (threadIdx.x == 0) {
        float recon_per_sample = partial_recon / batch_size;
        float kl_per_sample = partial_kl / batch_size;
        atomicAdd(loss_out, recon_per_sample + kl_per_sample);
        atomicAdd(recon_loss_out, recon_per_sample);
        atomicAdd(kl_loss_out, kl_per_sample);
    }
}

std::vector<torch::Tensor> cvae_forward_cuda(
    torch::Tensor x,
    torch::Tensor y,
    torch::Tensor enc_weight,
    torch::Tensor enc_bias,
    torch::Tensor enc_mu_weight,
    torch::Tensor enc_mu_bias,
    torch::Tensor enc_logvar_weight,
    torch::Tensor enc_logvar_bias,
    torch::Tensor dec_weight1,
    torch::Tensor dec_bias1,
    torch::Tensor dec_weight2,
    torch::Tensor dec_bias2,
    int hidden_dim,
    int latent_dim
) {
    int batch_size = x.size(0);
    int input_dim = x.size(1);
    int label_dim = y.size(1);

    auto options = torch::TensorOptions().device(x.device()).dtype(torch::kFloat32);
    auto loss_out = torch::zeros({}, options);
    auto recon_out = torch::zeros({}, options);
    auto kl_out = torch::zeros({}, options);

    int threads = 128;
    size_t smem_size = (hidden_dim + 3*latent_dim + hidden_dim + input_dim) * sizeof(float);

    unsigned long long seed = 123456789ULL;

    cvae_forward_kernel<<<batch_size, threads, smem_size>>>(
        x.data_ptr<float>(),
        y.data_ptr<float>(),
        enc_weight.data_ptr<float>(),
        enc_bias.data_ptr<float>(),
        enc_mu_weight.data_ptr<float>(),
        enc_mu_bias.data_ptr<float>(),
        enc_logvar_weight.data_ptr<float>(),
        enc_logvar_bias.data_ptr<float>(),
        dec_weight1.data_ptr<float>(),
        dec_bias1.data_ptr<float>(),
        dec_weight2.data_ptr<float>(),
        dec_bias2.data_ptr<float>(),
        loss_out.data_ptr<float>(),
        recon_out.data_ptr<float>(),
        kl_out.data_ptr<float>(),
        batch_size, input_dim, label_dim, hidden_dim, latent_dim, seed
    );

    return {loss_out, recon_out, kl_out};
}
'''

cpp_source = r'''
std::vector<torch::Tensor> cvae_forward_cuda(
    torch::Tensor x, torch::Tensor y,
    torch::Tensor enc_weight, torch::Tensor enc_bias,
    torch::Tensor enc_mu_weight, torch::Tensor enc_mu_bias,
    torch::Tensor enc_logvar_weight, torch::Tensor enc_logvar_bias,
    torch::Tensor dec_weight1, torch::Tensor dec_bias1,
    torch::Tensor dec_weight2, torch::Tensor dec_bias2,
    int hidden_dim, int latent_dim
);
'''

module = load_inline(
    name='cvae_cuda',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['cvae_forward_cuda'],
    verbose=False,
    extra_cuda_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.input_dim = input_dim
        self.label_dim = label_dim
        self.hidden_dim = hidden_dim
        self.latent_dim = latent_dim

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> tuple:
        # Extract weights
        enc_w = self.enc[0].weight
        enc_b = self.enc[0].bias
        mu_w = self.enc_mu.weight
        mu_b = self.enc_mu.bias
        logvar_w = self.enc_logvar.weight
        logvar_b = self.enc_logvar.bias
        dec_w1 = self.dec[0].weight
        dec_b1 = self.dec[0].bias
        dec_w2 = self.dec[2].weight
        dec_b2 = self.dec[2].bias

        outs = module.cvae_forward_cuda(
            x, y, enc_w, enc_b, mu_w, mu_b, logvar_w, logvar_b,
            dec_w1, dec_b1, dec_w2, dec_b2, self.hidden_dim, self.latent_dim
        )
        loss = outs[0]
        recon = outs[1]
        kl = outs[2]
        return loss, recon, kl