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

class VAE(nn.Module):
    def __init__(self,image_size=784,h_dim=400,z_dim=20):
        super(VAE,self).__init__()
        self.fc1=nn.Linear(image_size,h_dim)
        self.fc2=nn.Linear(h_dim,z_dim)
        self.fc3=nn.Linear(h_dim,z_dim)
        self.fc4=nn.Linear(z_dim,h_dim)
        self.fc5=nn.Linear(h_dim,image_size)
    def encoder(self,x):
        h=F.relu(self.fc1(x))
        return self.fc2(h),self.fc3(h)
    
    def decoder(self,z):
        h=F.relu(self.fc4(z))
        return F.sigmoid(self.fc5(h))

    def reparameterize(self,mu,log_var):
        std=torch.exp(log_var/2)
        eps=torch.randn_like(std)
        return mu+eps*std

    def forward(self,x):
        mu,log_val=self.encoder(x)
        z=self.reparameterize(mu,log_val)
        x_recon=self.decoder(z)

        return x_recon,mu,log_val