import torch
import torch.nn as nn

class ConditionalVAE(nn.Module):
    def __init__(self, x_dim, c_dim, latent_dim):
        super(ConditionalVAE, self).__init__()
        self.encoder_fc = nn.Sequential(
            nn.Linear(x_dim + c_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU()
        )
        self.fc_mu = nn.Linear(128, latent_dim)
        self.fc_logvar = nn.Linear(128, latent_dim)
        
        self.decoder_fc = nn.Sequential(
            nn.Linear(latent_dim + c_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, x_dim)
        )
        
    def encode(self, x, c):
        inputs = torch.cat([x, c], dim=1)
        h = self.encoder_fc(inputs)
        return self.fc_mu(h), self.fc_logvar(h)
        
    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
        
    def decode(self, z, c):
        inputs = torch.cat([z, c], dim=1)
        return self.decoder_fc(inputs)
        
    def forward(self, x, c):
        mu, logvar = self.encode(x, c)
        z = self.reparameterize(mu, logvar)
        x_recon = self.decode(z, c)
        return x_recon, mu, logvar