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, 22, 7, 54), 'torch.nn.Linear', 'nn.Linear', ({(7, 32, 7, 43): 'signal_size', (7, 45, 7, 53): '1024 * 4 * 4'}, {}), '(signal_size, 1024 * 4 * 4)', True, 'import torch.nn as nn\n'), ((18, 21, 18, 42), 'torch.nn.Sequential', 'nn.Sequential', ({(18, 35, 18, 41): '*convs'}, {}), '(*convs)', True, 'import torch.nn as nn\n'), ((35, 21, 35, 42), 'torch.nn.Sequential', 'nn.Sequential', ({(35, 35, 35, 41): '*convs'}, {}), '(*convs)', True, 'import torch.nn as nn\n'), ((36, 22, 36, 44), 'torch.nn.Linear', 'nn.Linear', ({(36, 32, 36, 40): '1024 * 1 * 1', (36, 42, 36, 43): '1'}, {}), '(1024 * 1 * 1, 1)', True, 'import torch.nn as nn\n'), ((15, 21, 15, 80), 'torch.nn.ConvTranspose2d', 'nn.ConvTranspose2d', (), '', True, 'import torch.nn as nn\n'), ((16, 21, 16, 30), 'torch.nn.Tanh', 'nn.Tanh', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((12, 25, 12, 84), 'torch.nn.ConvTranspose2d', 'nn.ConvTranspose2d', (), '', True, 'import torch.nn as nn\n'), ((13, 25, 13, 52), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(13, 40, 13, 51): 'channels[i]'}, {}), '(channels[i])', True, 'import torch.nn as nn\n'), ((14, 25, 14, 56), 'torch.nn.LeakyReLU', 'nn.LeakyReLU', (), '', True, 'import torch.nn as nn\n'), ((32, 25, 32, 86), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((33, 25, 33, 52), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', ({(33, 40, 33, 51): 'channels[i]'}, {}), '(channels[i])', True, 'import torch.nn as nn\n'), ((34, 25, 34, 56), 'torch.nn.LeakyReLU', 'nn.LeakyReLU', (), '', True, 'import torch.nn as nn\n')]