import torch.nn as nn
import torch
from utils import Flatten , Unflatten , weights_init , down_conv , up_conv

class Net(nn.Module):
    def __init__(self , num_layers , img_dim , in_chan , act_func , latent_vector_size):
        
        super(Net , self).__init__()
        assert act_func in ("ReLU" , "LeakyReLU") , "Activation function that can be used now are ReLU and LeakyReLU"
        assert img_dim % (2**(num_layers)) >= 0 , "Latent vector driven to 0, please increase image size or decreasenumber of layers"
        
        self.act_func = act_func
        self.in_chan = in_chan
        self.num_layers = num_layers
        self.latent_vector_size = latent_vector_size
        self.in_chan2 = self.in_chan
        self.encoder_net_layers = []
        self.decoder_net_layers = []
        
        self.out_chan = 2**5
        for _ in range(num_layers):
            
            self.encoder_net_layers.append(down_conv(self.in_chan , self.act_func , self.out_chan))
            self.in_chan = self.out_chan*2
            self.out_chan = self.out_chan*4
        self.encoder = nn.Sequential(*self.encoder_net_layers , 
                                     Flatten() , 
                                     nn.Linear(((self.out_chan//2)*((img_dim//(2 ** num_layers))**2)) , self.latent_vector_size*4) , 
                                     nn.ReLU(),
                                     nn.Linear(self.latent_vector_size*4 , self.latent_vector_size*2) , 
                                     nn.ReLU()
                                    )
        self.mu = nn.Linear(self.latent_vector_size*2 , self.latent_vector_size)
        self.logvar = nn.Linear(self.latent_vector_size*2 , self.latent_vector_size)
        self.out_chan2 = self.out_chan
        
        for _ in range(num_layers):
            self.decoder_net_layers.append(up_conv(self.out_chan2//2 , self.act_func , self.out_chan2//4))
            self.out_chan2 = self.out_chan2//4
        
        
        
        self.decoder = nn.Sequential(nn.Linear(self.latent_vector_size , self.latent_vector_size*4) , 
                                    nn.ReLU() ,  
                                    nn.Linear(self.latent_vector_size*4 , ((self.out_chan//2)*((img_dim//(2 ** num_layers))**2))) , 
                                    nn.ReLU() , 
                                    Unflatten(self.out_chan//2 , (img_dim//(2 ** num_layers)) , (img_dim//(2 ** num_layers)) ) , 
                                    *self.decoder_net_layers , 
                                    nn.ConvTranspose2d(self.out_chan2//2 , self.in_chan2 , 3 , 1 , 1))
        
    def encode(self , input_tensor):
        encoded_vector = self.encoder(input_tensor)
        mu , logvar = self.mu(encoded_vector) , self.logvar(encoded_vector)
        return mu , logvar
    
    def reparameterize(self , mu , logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        latent = mu + std*eps
        return latent
    
    def decode(self , latent):
        decoded_vector = self.decoder(latent)
        return decoded_vector
    
    def forward(self , input_tensor):
        mu , logvar = self.encode(input_tensor)
        latent_space = self.reparameterize(mu , logvar)
        return self.decode(latent_space) , mu , logvar
[((33, 18, 33, 80), 'torch.nn.Linear', 'nn.Linear', ({(33, 28, 33, 53): 'self.latent_vector_size * 2', (33, 56, 33, 79): 'self.latent_vector_size'}, {}), '(self.latent_vector_size * 2, self.latent_vector_size)', True, 'import torch.nn as nn\n'), ((34, 22, 34, 84), 'torch.nn.Linear', 'nn.Linear', ({(34, 32, 34, 57): 'self.latent_vector_size * 2', (34, 60, 34, 83): 'self.latent_vector_size'}, {}), '(self.latent_vector_size * 2, self.latent_vector_size)', True, 'import torch.nn as nn\n'), ((57, 14, 57, 37), 'torch.exp', 'torch.exp', ({(57, 24, 57, 36): '0.5 * logvar'}, {}), '(0.5 * logvar)', False, 'import torch\n'), ((58, 14, 58, 35), 'torch.randn_like', 'torch.randn_like', ({(58, 31, 58, 34): 'std'}, {}), '(std)', False, 'import torch\n'), ((27, 37, 27, 46), 'utils.Flatten', 'Flatten', ({}, {}), '()', False, 'from utils import Flatten, Unflatten, weights_init, down_conv, up_conv\n'), ((28, 37, 28, 130), 'torch.nn.Linear', 'nn.Linear', ({(28, 48, 28, 100): 'self.out_chan // 2 * (img_dim // 2 ** num_layers) ** 2', (28, 104, 28, 129): 'self.latent_vector_size * 4'}, {}), '(self.out_chan // 2 * (img_dim // 2 ** num_layers) ** 2, self.\n    latent_vector_size * 4)', True, 'import torch.nn as nn\n'), ((29, 37, 29, 46), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((30, 37, 30, 101), 'torch.nn.Linear', 'nn.Linear', ({(30, 47, 30, 72): 'self.latent_vector_size * 4', (30, 75, 30, 100): 'self.latent_vector_size * 2'}, {}), '(self.latent_vector_size * 4, self.latent_vector_size * 2)', True, 'import torch.nn as nn\n'), ((31, 37, 31, 46), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((43, 37, 43, 99), 'torch.nn.Linear', 'nn.Linear', ({(43, 47, 43, 70): 'self.latent_vector_size', (43, 73, 43, 98): 'self.latent_vector_size * 4'}, {}), '(self.latent_vector_size, self.latent_vector_size * 4)', True, 'import torch.nn as nn\n'), ((44, 36, 44, 45), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((45, 36, 45, 129), 'torch.nn.Linear', 'nn.Linear', ({(45, 46, 45, 71): 'self.latent_vector_size * 4', (45, 75, 45, 127): 'self.out_chan // 2 * (img_dim // 2 ** num_layers) ** 2'}, {}), '(self.latent_vector_size * 4, self.out_chan // 2 * (img_dim // 2 **\n    num_layers) ** 2)', True, 'import torch.nn as nn\n'), ((46, 36, 46, 45), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((47, 36, 47, 126), 'utils.Unflatten', 'Unflatten', ({(47, 46, 47, 62): 'self.out_chan // 2', (47, 66, 47, 92): 'img_dim // 2 ** num_layers', (47, 97, 47, 123): 'img_dim // 2 ** num_layers'}, {}), '(self.out_chan // 2, img_dim // 2 ** num_layers, img_dim // 2 **\n    num_layers)', False, 'from utils import Flatten, Unflatten, weights_init, down_conv, up_conv\n'), ((49, 36, 49, 101), 'torch.nn.ConvTranspose2d', 'nn.ConvTranspose2d', ({(49, 55, 49, 72): 'self.out_chan2 // 2', (49, 75, 49, 88): 'self.in_chan2', (49, 91, 49, 92): '3', (49, 95, 49, 96): '1', (49, 99, 49, 100): '1'}, {}), '(self.out_chan2 // 2, self.in_chan2, 3, 1, 1)', True, 'import torch.nn as nn\n'), ((23, 43, 23, 98), 'utils.down_conv', 'down_conv', ({(23, 53, 23, 65): 'self.in_chan', (23, 68, 23, 81): 'self.act_func', (23, 84, 23, 97): 'self.out_chan'}, {}), '(self.in_chan, self.act_func, self.out_chan)', False, 'from utils import Flatten, Unflatten, weights_init, down_conv, up_conv\n'), ((38, 43, 38, 105), 'utils.up_conv', 'up_conv', ({(38, 51, 38, 68): '(self.out_chan2 // 2)', (38, 71, 38, 84): 'self.act_func', (38, 87, 38, 104): '(self.out_chan2 // 4)'}, {}), '(self.out_chan2 // 2, self.act_func, self.out_chan2 // 4)', False, 'from utils import Flatten, Unflatten, weights_init, down_conv, up_conv\n')]