# Training settings

parser = argparse.ArgumentParser(description='pix2pix-pytorch-implementation')

parser.add_argument('--dataset', required=True, help='facades')

parser.add_argument('--batch_size', type=int, default=1, help='training batch size')

parser.add_argument('--test_batch_size', type=int, default=1, help='testing batch size')

parser.add_argument('--direction', type=str, default='b2a', help='a2b or b2a')

parser.add_argument('--input_nc', type=int, default=3, help='input image channels')

parser.add_argument('--output_nc', type=int, default=3, help='output image channels')

parser.add_argument('--ngf', type=int, default=64, help='generator filters in first conv layer')

parser.add_argument('--ndf', type=int, default=64, help='discriminator filters in first conv layer')

parser.add_argument('--epoch_count', type=int, default=1, help='the starting epoch count')

parser.add_argument('--niter', type=int, default=100, help='# of iter at starting learning rate')

parser.add_argument('--niter_decay', type=int, default=100, help='# of iter to linearly decay learning rate to zero')

parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate for adam')

parser.add_argument('--lr_policy', type=str, default='lambda', help='learning rate policy: lambda|step|plateau|cosine')

parser.add_argument('--lr_decay_iters', type=int, default=50, help='multiply by a gamma every lr_decay_iters iterations')

parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5')

parser.add_argument('--cuda', action='store_true', help='use cuda?')

parser.add_argument('--threads', type=int, default=4, help='number of threads for data loader to use')

parser.add_argument('--seed', type=int, default=123, help='random seed to use. Default=123')

parser.add_argument('--lamb', type=int, default=10, help='weight on L1 term in objective')

opt = parser.parse_args()

print('===> Loading datasets')

root_path = "dataset/"

train_set = get_training_set(root_path + opt.dataset, opt.direction)

test_set = get_test_set(root_path + opt.dataset, opt.direction)

training_data_loader = DataLoader(dataset=train_set, num_workers=opt.threads, batch_size=opt.batch_size, shuffle=True)

testing_data_loader = DataLoader(dataset=test_set, num_workers=opt.threads, batch_size=opt.test_batch_size, shuffle=False)

print('===> Building models')

net_g = define_G(opt.input_nc, opt.output_nc, opt.ngf, 'batch', False, 'normal', 0.02, gpu_id=device)

net_d = define_D(opt.input_nc + opt.output_nc, opt.ndf, 'basic', gpu_id=device)

criterionGAN = GANLoss().to(device)

criterionL1 = nn.L1Loss().to(device)

criterionMSE = nn.MSELoss().to(device)

# setup optimizer

optimizer_g = optim.Adam(net_g.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))

optimizer_d = optim.Adam(net_d.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))

net_g_scheduler = get_scheduler(optimizer_g, opt)

net_d_scheduler = get_scheduler(optimizer_d, opt)

######################

        # (1) Update D network

        ######################

        optimizer_d.zero_grad()

        # train with fake

        fake_ab = torch.cat((real_a, fake_b), 1)

        pred_fake = net_d.forward(fake_ab.detach())

        loss_d_fake = criterionGAN(pred_fake, False)

        # train with real

        real_ab = torch.cat((real_a, real_b), 1)

        pred_real = net_d.forward(real_ab)

        loss_d_real = criterionGAN(pred_real, True)

        # Combined D loss

        loss_d = (loss_d_fake + loss_d_real) * 0.5

        loss_d.backward()

        optimizer_d.step()

######################

        # (2) Update G network

        ######################

        optimizer_g.zero_grad()

        # First, G(A) should fake the discriminator

        fake_ab = torch.cat((real_a, fake_b), 1)

        pred_fake = net_d.forward(fake_ab)

        loss_g_gan = criterionGAN(pred_fake, True)

        # Second, G(A) = B

        loss_g_l1 = criterionL1(fake_b, real_b) * opt.lamb

        loss_g = loss_g_gan + loss_g_l1

        loss_g.backward()

        optimizer_g.step()

update_learning_rate(net_g_scheduler, optimizer_g)

update_learning_rate(net_d_scheduler, optimizer_d)