from torch import nn


def get_fc_discriminator(num_classes, ndf=64):
    return nn.Sequential(
        nn.Conv2d(num_classes, ndf, kernel_size=4, stride=2, padding=1),
        nn.LeakyReLU(negative_slope=0.2, inplace=True),
        nn.Conv2d(ndf, ndf * 2, kernel_size=4, stride=2, padding=1),
        nn.LeakyReLU(negative_slope=0.2, inplace=True),
        nn.Conv2d(ndf * 2, ndf * 4, kernel_size=4, stride=2, padding=1),
        nn.LeakyReLU(negative_slope=0.2, inplace=True),
        nn.Conv2d(ndf * 4, ndf * 8, kernel_size=4, stride=2, padding=1),
        nn.LeakyReLU(negative_slope=0.2, inplace=True),
        nn.Conv2d(ndf * 8, 1, kernel_size=4, stride=2, padding=1),
    )


# def get_fe_discriminator(num_classes, ndf=64): # 256-128-64-32-16
#     return nn.Sequential(
#         nn.Conv2d(num_classes, ndf * 4, kernel_size=4, stride=2, padding=1),
#         nn.LeakyReLU(negative_slope=0.2, inplace=True),
#         nn.Conv2d(ndf * 4, ndf * 2, kernel_size=4, stride=2, padding=1),
#         nn.LeakyReLU(negative_slope=0.2, inplace=True),
#         nn.Conv2d(ndf * 2, ndf, kernel_size=2, stride=2, padding=0),
#         nn.LeakyReLU(negative_slope=0.2, inplace=True),
#         # nn.Conv2d(ndf * 4, ndf * 8, kernel_size=4, stride=2, padding=1),
#         # nn.LeakyReLU(negative_slope=0.2, inplace=True),
#         nn.Conv2d(ndf, 1, kernel_size=2, stride=2, padding=0),
#     )

# def get_fe_discriminator(num_classes, ndf=64):
#     return nn.Sequential(
#         nn.Conv2d(num_classes, ndf, kernel_size=4, stride=2, padding=1),
#         nn.LeakyReLU(negative_slope=0.2, inplace=True),
#         nn.Conv2d(ndf, ndf * 2, kernel_size=4, stride=2, padding=1),
#         nn.LeakyReLU(negative_slope=0.2, inplace=True),
#         nn.Conv2d(ndf * 2, ndf * 4, kernel_size=4, stride=2, padding=1),
#         nn.LeakyReLU(negative_slope=0.2, inplace=True),
#         # nn.Conv2d(ndf * 4, ndf * 8, kernel_size=4, stride=2, padding=1),
#         # nn.LeakyReLU(negative_slope=0.2, inplace=True),
#         nn.Conv2d(ndf * 4, 1, kernel_size=1, stride=1, padding=0),
#     )

def get_fe_discriminator(num_classes, ndf=64):  # H/8,H/8,(1024 -> 256 -> 128 -> 64  -> 1)
    return nn.Sequential(
        nn.Conv2d(num_classes, ndf * 4, kernel_size=1, stride=1, padding=0),
        # x=self.dropout(x)
        nn.LeakyReLU(negative_slope=0.2, inplace=True),
        nn.Conv2d(ndf * 4, ndf * 2, kernel_size=1, stride=1, padding=0),
        # x=self.dropout(x)
        nn.LeakyReLU(negative_slope=0.2, inplace=True),
        nn.Conv2d(ndf * 2, ndf, kernel_size=1, stride=1, padding=0),
        # x=self.dropout(x)
        nn.LeakyReLU(negative_slope=0.2, inplace=True),
        # nn.Conv2d(ndf * 4, ndf * 8, kernel_size=4, stride=2, padding=1),
        # nn.LeakyReLU(negative_slope=0.2, inplace=True),
        nn.Conv2d(ndf, 1, kernel_size=1, stride=1, padding=0),
    )
[((104, 167), 'torch.nn.Conv2d', 'nn.Conv2d', (['num_classes', 'ndf'], {'kernel_size': '(4)', 'stride': '(2)', 'padding': '(1)'}), '(num_classes, ndf, kernel_size=4, stride=2, padding=1)\n', (113, 167), False, 'from torch import nn\n'), ((177, 223), 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ([], {'negative_slope': '(0.2)', 'inplace': '(True)'}), '(negative_slope=0.2, inplace=True)\n', (189, 223), False, 'from torch import nn\n'), ((233, 292), 'torch.nn.Conv2d', 'nn.Conv2d', (['ndf', '(ndf * 2)'], {'kernel_size': '(4)', 'stride': '(2)', 'padding': '(1)'}), '(ndf, ndf * 2, kernel_size=4, stride=2, padding=1)\n', (242, 292), False, 'from torch import nn\n'), ((302, 348), 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ([], {'negative_slope': '(0.2)', 'inplace': '(True)'}), '(negative_slope=0.2, inplace=True)\n', (314, 348), False, 'from torch import nn\n'), ((358, 421), 'torch.nn.Conv2d', 'nn.Conv2d', (['(ndf * 2)', '(ndf * 4)'], {'kernel_size': '(4)', 'stride': '(2)', 'padding': '(1)'}), '(ndf * 2, ndf * 4, kernel_size=4, stride=2, padding=1)\n', (367, 421), False, 'from torch import nn\n'), ((431, 477), 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ([], {'negative_slope': '(0.2)', 'inplace': '(True)'}), '(negative_slope=0.2, inplace=True)\n', (443, 477), False, 'from torch import nn\n'), ((487, 550), 'torch.nn.Conv2d', 'nn.Conv2d', (['(ndf * 4)', '(ndf * 8)'], {'kernel_size': '(4)', 'stride': '(2)', 'padding': '(1)'}), '(ndf * 4, ndf * 8, kernel_size=4, stride=2, padding=1)\n', (496, 550), False, 'from torch import nn\n'), ((560, 606), 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ([], {'negative_slope': '(0.2)', 'inplace': '(True)'}), '(negative_slope=0.2, inplace=True)\n', (572, 606), False, 'from torch import nn\n'), ((616, 673), 'torch.nn.Conv2d', 'nn.Conv2d', (['(ndf * 8)', '(1)'], {'kernel_size': '(4)', 'stride': '(2)', 'padding': '(1)'}), '(ndf * 8, 1, kernel_size=4, stride=2, padding=1)\n', (625, 673), False, 'from torch import nn\n'), ((2201, 2268), 'torch.nn.Conv2d', 'nn.Conv2d', (['num_classes', '(ndf * 4)'], {'kernel_size': '(1)', 'stride': '(1)', 'padding': '(0)'}), '(num_classes, ndf * 4, kernel_size=1, stride=1, padding=0)\n', (2210, 2268), False, 'from torch import nn\n'), ((2306, 2352), 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ([], {'negative_slope': '(0.2)', 'inplace': '(True)'}), '(negative_slope=0.2, inplace=True)\n', (2318, 2352), False, 'from torch import nn\n'), ((2362, 2425), 'torch.nn.Conv2d', 'nn.Conv2d', (['(ndf * 4)', '(ndf * 2)'], {'kernel_size': '(1)', 'stride': '(1)', 'padding': '(0)'}), '(ndf * 4, ndf * 2, kernel_size=1, stride=1, padding=0)\n', (2371, 2425), False, 'from torch import nn\n'), ((2463, 2509), 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ([], {'negative_slope': '(0.2)', 'inplace': '(True)'}), '(negative_slope=0.2, inplace=True)\n', (2475, 2509), False, 'from torch import nn\n'), ((2519, 2578), 'torch.nn.Conv2d', 'nn.Conv2d', (['(ndf * 2)', 'ndf'], {'kernel_size': '(1)', 'stride': '(1)', 'padding': '(0)'}), '(ndf * 2, ndf, kernel_size=1, stride=1, padding=0)\n', (2528, 2578), False, 'from torch import nn\n'), ((2616, 2662), 'torch.nn.LeakyReLU', 'nn.LeakyReLU', ([], {'negative_slope': '(0.2)', 'inplace': '(True)'}), '(negative_slope=0.2, inplace=True)\n', (2628, 2662), False, 'from torch import nn\n'), ((2805, 2858), 'torch.nn.Conv2d', 'nn.Conv2d', (['ndf', '(1)'], {'kernel_size': '(1)', 'stride': '(1)', 'padding': '(0)'}), '(ndf, 1, kernel_size=1, stride=1, padding=0)\n', (2814, 2858), False, 'from torch import nn\n')]