nll = nn.NLLLoss().cuda()
mse = nn.MSELoss().cuda()
bce = nn.BCELoss().cuda()

# plt.figure(0, figsize = (32, 32))
for epoch in range(100):
    # pb = tqdm_notebook(total = train_loader.dataset.data_tensor.size()[0])
    # pb.set_description('Epoch ' + str(epoch + 1))
    for i, (data, targets) in enumerate(train_loader):
        ones = Variable(torch.ones(data.size()[0], 1)).cuda()
        zeros = Variable(torch.zeros(data.size()[0], 1)).cuda()

        z_dict = self.get_z(data.size()[0])
        z = torch.cat([z_dict[k] for k in z_dict.keys()], dim=1)

        data = Variable(data.float().cuda(async=True)) / 255.0
        targets = Variable(targets.float().cuda(async=True))

        # Forward pass on real MNIST
        out_dis, hid = self.dis(data)
        c1 = F.log_softmax(self.Q_cat(hid))
        loss_dis = mse(out_dis, ones) - torch.sum(targets * c1) / (torch.sum(targets) + 1e-3)  # Loss for real MNIST

        # Forward pass on generated MNIST
        out_gen = self.gen(z)
        out_dis, hid = self.dis(out_gen)

        # Loss for generated MNIST
        loss_dis = loss_dis + mse(out_dis, zeros)
        # loss_dis = loss_dis


        # Zero gradient buffers for gen and Q_cat and backward pass
        self.dis.zero_grad()
        self.Q_cat.zero_grad()
        loss_dis.backward(retain_graph=True)  # We need PyTorch to retain the graph buffers so we can run backward again later
        self.d_optim.step()  # Apply the discriminator's update now since we have to delete its gradients later

        # And backward pass and loss for generator and update
        self.gen.zero_grad()
        loss_gen = mse(out_dis, ones)
        loss_gen.backward(retain_graph=True)
        self.dis.zero_grad()
        # Don't want the gradients of the generator's objective in the discriminator

        # Forward pass and loss for latent codes
        # loss_q = 0

        c1 = F.log_softmax(self.Q_cat(hid))
        loss_q = nll(c1, torch.max(z_dict['cat'], dim=1)[1])

        if self.c2_len:
            c2 = self.Q_con(hid)
            loss_q += 0.5 * mse(c2, z_dict[
                'con'])  # Multiply by 0.5 as we treat targets as Gaussian (and there's a coefficient of 0.5 when we take logs)
            self.Q_con.zero_grad()  # Zero gradient buffers before the backward pass
        if self.c3_len:
            c3 = F.sigmoid(self.Q_bin(hid))
            loss_q += bce(c3, z_dict['bin'])
            self.Q_bin.zero_grad()  # Zero gradient buffers before the backward pass
        # Backward pass for latent code objective
        loss_q.backward()

        # Do the updates for everything
        self.d_optim.step()
        self.g_optim.step()
        self.qcat_optim.step()

        if self.c2_len:
            self.qcon_optim.step()
        if self.c3_len:
            self.qbin_optim.step()

            # pb.update(data.size()[0])
            # pb.set_postfix(loss_dis = loss_dis.cpu().data.numpy()[0], loss_gen = loss_gen.cpu().data.numpy()[0], loss_q = loss_q.cpu().data.numpy()[0])
    print("epoch %d discriminator loss for generated MNIST %f") % (epoch, loss_dis.cpu().data.numpy()[0])
    print("epoch %d generation loss %f") % (epoch, loss_gen.cpu().data.numpy()[0])
    print("epoch %d latent loss %f") % (epoch, loss_q.cpu().data.numpy()[0])
    # pb.close()
    # plt.subplot(10, 10, epoch + 1)
    # bp()
    x = np.squeeze(np.transpose(out_gen.cpu().data.numpy(), (0, 2, 3, 1)))
    true_x = np.squeeze(np.transpose(data.cpu().data.numpy(), (0, 2, 3, 1)))
    merged_x = merge(x[:64], [8, 8])
    merged_true = merge(true_x[:64], [8, 8])
    # bp()
    scipy.misc.imsave("./samples/gn_" + str(epoch) + ".png", merged_x)
    scipy.misc.imsave("./samples/gr_" + str(epoch) + ".png", merged_true)