import torch
#Author: @meetdoshi
#Reference: https://github.com/Lornatang/CycleGAN-PyTorch/blob/master/cyclegan_pytorch/models.py
class Discriminator(torch.nn.Module):
    '''
    PatchGAN Discriminator with 70x70 overlapping image patches
    '''
    def __init__(self):
        super(Discriminator,self).__init__()
        self.model = torch.nn.Sequential(
            torch.nn.Conv2d(3, 64, 4, stride=2, padding=1),
            torch.nn.LeakyReLU(0.2, True),
            torch.nn.Conv2d(64, 128, 4,stride=2, padding=1),
            torch.nn.InstanceNorm2d(128),
            torch.nn.LeakyReLU(0.2, True),
            torch.nn.Conv2d(128, 256, 4, stride=2, padding=1),
            torch.nn.InstanceNorm2d(256),
            torch.nn.LeakyReLU(0.2, True),
            torch.nn.Conv2d(256, 512, 4,padding=1),
            torch.nn.InstanceNorm2d(512),
            torch.nn.LeakyReLU(0.2, True),
            torch.nn.Conv2d(512, 1, 4, padding=1),
        )

    def forward(self, x):
        x = self.model(x)
        x = torch.nn.functional.avg_pool2d(x,x.size()[2:])
        x = torch.flatten(x,1)
        return x