import torch
import torch.nn as nn

class Generator(nn.Module):
    def __init__(self, signal_size, out_channels=3):
        super(Generator, self).__init__()
        self.linear = nn.Linear(signal_size, 1024*4*4)
        
        convs = []
        channels = [1024, 512, 256, 128]
        for i in range(1, len(channels)):
            convs.append(nn.ConvTranspose2d(channels[i-1], channels[i], 2, stride=2))
            convs.append(nn.BatchNorm2d(channels[i]))
            convs.append(nn.LeakyReLU(0.2, inplace=True))
        convs.append(nn.ConvTranspose2d(channels[-1], out_channels, 2, stride=2))
        convs.append(nn.Tanh())
        
        self.convs = nn.Sequential(*convs)
        
    def forward(self, x):
        x = self.linear(x)
        x = x.view(x.size(0), 1024, 4, 4)
        x = self.convs(x)
        return x

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        channels = [3, 32, 64, 128, 256, 512, 1024]
        convs = []
        for i in range(1, len(channels)):
            convs.append(nn.Conv2d(channels[i-1], channels[i], 3, padding=1, stride=2))
            convs.append(nn.BatchNorm2d(channels[i]))
            convs.append(nn.LeakyReLU(0.2, inplace=True))
        self.convs = nn.Sequential(*convs)
        self.linear = nn.Linear(1024*1*1, 1)
    
    def forward(self, x):
        x = self.convs(x)
        x = x.view(x.size(0), -1)      
        x = self.linear(x)
        # x = torch.sigmoid(x)
        return x

    

[(7, 'torch.nn.Linear', 'nn.Linear', (['signal_size', '(1024 * 4 * 4)'], {}), True, 'import torch.nn as nn\n'), (18, 'torch.nn.Sequential', 'nn.Sequential', (['*convs'], {}), True, 'import torch.nn as nn\n'), (35, 'torch.nn.Sequential', 'nn.Sequential', (['*convs'], {}), True, 'import torch.nn as nn\n'), (36, 'torch.nn.Linear', 'nn.Linear', (['(1024 * 1 * 1)', '(1)'], {}), True, 'import torch.nn as nn\n'), (15, 'torch.nn.ConvTranspose2d', 'nn.ConvTranspose2d', (['channels[-1]', 'out_channels', '(2)'], {'stride': '(2)'}), True, 'import torch.nn as nn\n'), (16, 'torch.nn.Tanh', 'nn.Tanh', ([], {}), True, 'import torch.nn as nn\n'), (12, 'torch.nn.ConvTranspose2d', 'nn.ConvTranspose2d', (['channels[i - 1]', 'channels[i]', '(2)'], {'stride': '(2)'}), True, 'import torch.nn as nn\n'), (13, 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['channels[i]'], {}), True, 'import torch.nn as nn\n'), (14, 'torch.nn.LeakyReLU', 'nn.LeakyReLU', (['(0.2)'], {'inplace': '(True)'}), True, 'import torch.nn as nn\n'), (32, 'torch.nn.Conv2d', 'nn.Conv2d', (['channels[i - 1]', 'channels[i]', '(3)'], {'padding': '(1)', 'stride': '(2)'}), True, 'import torch.nn as nn\n'), (33, 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['channels[i]'], {}), True, 'import torch.nn as nn\n'), (34, 'torch.nn.LeakyReLU', 'nn.LeakyReLU', (['(0.2)'], {'inplace': '(True)'}), True, 'import torch.nn as nn\n')]