from __future__ import print_function
from six.moves import range

import torchvision.transforms as transforms
import torch.backends.cudnn as cudnn
import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.optim as optim
import torchvision.utils as vutils
import numpy as np
import os
import time
from PIL import Image, ImageFont, ImageDraw
from copy import deepcopy

from miscc.config import cfg
from miscc.utils import mkdir_p

from CaptionDatasets import *

from tensorboard import summary
from tensorboard import FileWriter

from model import G_NET, D_NET64, D_NET128, D_NET256, D_NET512, D_NET1024, INCEPTION_V3



# ################## Shared functions ###################
def compute_mean_covariance(img):
    batch_size = img.size(0)
    channel_num = img.size(1)
    height = img.size(2)
    width = img.size(3)
    num_pixels = height * width

    # batch_size * channel_num * 1 * 1
    mu = img.mean(2, keepdim=True).mean(3, keepdim=True)

    # batch_size * channel_num * num_pixels
    img_hat = img - mu.expand_as(img)
    img_hat = img_hat.view(batch_size, channel_num, num_pixels)
    # batch_size * num_pixels * channel_num
    img_hat_transpose = img_hat.transpose(1, 2)
    # batch_size * channel_num * channel_num
    covariance = torch.bmm(img_hat, img_hat_transpose)
    covariance = covariance / num_pixels

    return mu, covariance


def KL_loss(mu, logvar):
    # -0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    KLD_element = mu.pow(2).add_(logvar.exp()).mul_(-1).add_(1).add_(logvar)
    KLD = torch.mean(KLD_element).mul_(-0.5)
    return KLD


def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        nn.init.orthogonal(m.weight.data, 1.0)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)
    elif classname.find('Linear') != -1:
        nn.init.orthogonal(m.weight.data, 1.0)
        if m.bias is not None:
            m.bias.data.fill_(0.0)


def load_params(model, new_param):
    for p, new_p in zip(model.parameters(), new_param):
        p.data.copy_(new_p)


def copy_G_params(model):
    flatten = deepcopy(list(p.data for p in model.parameters()))
    return flatten


def compute_inception_score(predictions, num_splits=1):
    # print('predictions', predictions.shape)
    scores = []
    for i in range(num_splits):
        istart = i * predictions.shape[0] // num_splits
        iend = (i + 1) * predictions.shape[0] // num_splits
        part = predictions[istart:iend, :]
        kl = part * \
            (np.log(part) - np.log(np.expand_dims(np.mean(part, 0), 0)))
        kl = np.mean(np.sum(kl, 1))
        scores.append(np.exp(kl))
    return np.mean(scores), np.std(scores)


def negative_log_posterior_probability(predictions, num_splits=1):
    # print('predictions', predictions.shape)
    scores = []
    for i in range(num_splits):
        istart = i * predictions.shape[0] // num_splits
        iend = (i + 1) * predictions.shape[0] // num_splits
        part = predictions[istart:iend, :]
        result = -1. * np.log(np.max(part, 1))
        result = np.mean(result)
        scores.append(result)
    return np.mean(scores), np.std(scores)


def load_network(gpus):
    netG = G_NET()
    netG.apply(weights_init)
    netG = torch.nn.DataParallel(netG, device_ids=gpus)
    print(netG)

    netsD = []
    if cfg.TREE.BRANCH_NUM > 0:
        netsD.append(D_NET64())
    if cfg.TREE.BRANCH_NUM > 1:
        netsD.append(D_NET128())
    if cfg.TREE.BRANCH_NUM > 2:
        netsD.append(D_NET256())
    if cfg.TREE.BRANCH_NUM > 3:
        netsD.append(D_NET512())
    if cfg.TREE.BRANCH_NUM > 4:
        netsD.append(D_NET1024())
    # TODO: if cfg.TREE.BRANCH_NUM > 5:

    for i in range(len(netsD)):
        netsD[i].apply(weights_init)
        netsD[i] = torch.nn.DataParallel(netsD[i], device_ids=gpus)
        # print(netsD[i])
    print('# of netsD', len(netsD))

    count = 0
    if cfg.TRAIN.NET_G != '':
        state_dict = torch.load(cfg.TRAIN.NET_G)
        netG.load_state_dict(state_dict)
        print('Load ', cfg.TRAIN.NET_G)

        istart = cfg.TRAIN.NET_G.rfind('_') + 1
        iend = cfg.TRAIN.NET_G.rfind('.')
        count = cfg.TRAIN.NET_G[istart:iend]
        count = int(count) + 1

    if cfg.TRAIN.NET_D != '':
        for i in range(len(netsD)):
            print('Load %s_%d.pth' % (cfg.TRAIN.NET_D, i))
            state_dict = torch.load('%s%d.pth' % (cfg.TRAIN.NET_D, i))
            netsD[i].load_state_dict(state_dict)

    inception_model = INCEPTION_V3()

    if cfg.CUDA:
        netG.cuda()
        for i in range(len(netsD)):
            netsD[i].cuda()
        inception_model = inception_model.cuda()
    inception_model.eval()

    return netG, netsD, len(netsD), inception_model, count


def define_optimizers(netG, netsD):
    optimizersD = []
    num_Ds = len(netsD)
    for i in range(num_Ds):
        opt = optim.Adam(netsD[i].parameters(),
                         lr=cfg.TRAIN.DISCRIMINATOR_LR,
                         betas=(0.5, 0.999))
        optimizersD.append(opt)

    # G_opt_paras = []
    # for p in netG.parameters():
    #     if p.requires_grad:
    #         G_opt_paras.append(p)
    optimizerG = optim.Adam(netG.parameters(),
                            lr=cfg.TRAIN.GENERATOR_LR,
                            betas=(0.5, 0.999))
    return optimizerG, optimizersD


def save_model(netG, avg_param_G, netsD, epoch, model_dir):
    load_params(netG, avg_param_G)
    torch.save(
        netG.state_dict(),
        '%s/netG_%d.pth' % (model_dir, epoch))
    for i in range(len(netsD)):
        netD = netsD[i]
        torch.save(
            netD.state_dict(),
            '%s/netD%d.pth' % (model_dir, i))
    print('Save G/Ds models.')

def save_real(imgs_tcpu, image_dir):
    num = cfg.TRAIN.VIS_COUNT

    # The range of real_img (i.e., self.imgs_tcpu[i][0:num])
    # is changed to [0, 1] by function vutils.save_image
    real_img = imgs_tcpu[-1][0:num]
    vutils.save_image(
        real_img, '%s/real_samples.png' % (image_dir),
        normalize=True)
    real_img_set = vutils.make_grid(real_img).numpy()
    real_img_set = np.transpose(real_img_set, (1, 2, 0))
    real_img_set = real_img_set * 255
    real_img_set = real_img_set.astype(np.uint8)
    sup_real_img = summary.image('real_img', real_img_set)


def save_img_results(imgs_tcpu, fake_imgs, num_imgs,
                     count, image_dir, summary_writer):
    num = cfg.TRAIN.VIS_COUNT

    # The range of real_img (i.e., self.imgs_tcpu[i][0:num])
    # is changed to [0, 1] by function vutils.save_image
    real_img = imgs_tcpu[-1][0:num]
    vutils.save_image(
        real_img, '%s/real_samples.png' % (image_dir),
        normalize=True)
    real_img_set = vutils.make_grid(real_img).numpy()
    real_img_set = np.transpose(real_img_set, (1, 2, 0))
    real_img_set = real_img_set * 255
    real_img_set = real_img_set.astype(np.uint8)
    sup_real_img = summary.image('real_img', real_img_set)
    summary_writer.add_summary(sup_real_img, count)

    for i in range(num_imgs):
        fake_img = fake_imgs[i][0:num]
        # The range of fake_img.data (i.e., self.fake_imgs[i][0:num])
        # is still [-1. 1]...
        vutils.save_image(
            fake_img.data, '%s/count_%09d_fake_samples_%d.png' %
            (image_dir, count, i), normalize=True)

        fake_img_set = vutils.make_grid(fake_img.data).cpu().numpy()

        fake_img_set = np.transpose(fake_img_set, (1, 2, 0))
        fake_img_set = (fake_img_set + 1) * 255 / 2
        fake_img_set = fake_img_set.astype(np.uint8)

        sup_fake_img = summary.image('fake_img%d' % i, fake_img_set)
        summary_writer.add_summary(sup_fake_img, count)
        summary_writer.flush()


# ################# Text to image task############################ #
class condGANTrainer(object):
    def __init__(self, output_dir, data_loader, imsize):
        if cfg.TRAIN.FLAG:
            self.model_dir = os.path.join(output_dir, 'Model')
            self.image_dir = os.path.join(output_dir, 'Image')
            self.log_dir = os.path.join(output_dir, 'Log')
            mkdir_p(self.model_dir)
            mkdir_p(self.image_dir)
            mkdir_p(self.log_dir)
            self.summary_writer = FileWriter(self.log_dir)

        s_gpus = cfg.GPU_ID.split(',')
        self.gpus = [int(ix) for ix in s_gpus]
        self.num_gpus = len(self.gpus)
        torch.cuda.set_device(self.gpus[0])
        cudnn.benchmark = True

        self.batch_size = cfg.TRAIN.BATCH_SIZE * self.num_gpus
        self.max_epoch = cfg.TRAIN.MAX_EPOCH
        self.snapshot_interval = cfg.TRAIN.SNAPSHOT_INTERVAL

        self.data_loader = data_loader
        self.num_batches = len(self.data_loader)

    def prepare_data(self, data):
        imgs, w_imgs, t_embedding, _ = data

        real_vimgs, wrong_vimgs = [], []
        if cfg.CUDA:
            vembedding = Variable(t_embedding).cuda()
        else:
            vembedding = Variable(t_embedding)
        for i in range(self.num_Ds):
            if cfg.CUDA:
                real_vimgs.append(Variable(imgs[i]).cuda())
                wrong_vimgs.append(Variable(w_imgs[i]).cuda())
            else:
                real_vimgs.append(Variable(imgs[i]))
                wrong_vimgs.append(Variable(w_imgs[i]))
        return imgs, real_vimgs, wrong_vimgs, vembedding

    def train_Dnet(self, idx, count):
        flag = count % 100
        batch_size = self.real_imgs[0].size(0)
        criterion, mu = self.criterion, self.mu

        netD, optD = self.netsD[idx], self.optimizersD[idx]
        real_imgs = self.real_imgs[idx]
        wrong_imgs = self.wrong_imgs[idx]
        fake_imgs = self.fake_imgs[idx]
        #
        netD.zero_grad()
        # Forward
        real_labels = self.real_labels[:batch_size]
        fake_labels = self.fake_labels[:batch_size]
        # for real
        real_logits = netD(real_imgs, mu.detach())
        wrong_logits = netD(wrong_imgs, mu.detach())
        fake_logits = netD(fake_imgs.detach(), mu.detach())
        #
        errD_real = criterion(real_logits[0], real_labels)
        errD_wrong = criterion(wrong_logits[0], fake_labels)
        errD_fake = criterion(fake_logits[0], fake_labels)
        if len(real_logits) > 1 and cfg.TRAIN.COEFF.UNCOND_LOSS > 0:
            errD_real_uncond = cfg.TRAIN.COEFF.UNCOND_LOSS * \
                criterion(real_logits[1], real_labels)
            errD_wrong_uncond = cfg.TRAIN.COEFF.UNCOND_LOSS * \
                criterion(wrong_logits[1], real_labels)
            errD_fake_uncond = cfg.TRAIN.COEFF.UNCOND_LOSS * \
                criterion(fake_logits[1], fake_labels)
            #
            errD_real = errD_real + errD_real_uncond
            errD_wrong = errD_wrong + errD_wrong_uncond
            errD_fake = errD_fake + errD_fake_uncond
            #
            errD = errD_real + errD_wrong + errD_fake
        else:
            errD = errD_real + 0.5 * (errD_wrong + errD_fake)
        # backward
        errD.backward()
        # update parameters
        optD.step()
        # log
        if flag == 0:
            summary_D = summary.scalar('D_loss%d' % idx, errD.item())
            self.summary_writer.add_summary(summary_D, count)
        return errD

    def train_Gnet(self, count):
        self.netG.zero_grad()
        errG_total = 0
        flag = count % 100
        batch_size = self.real_imgs[0].size(0)
        criterion, mu, logvar = self.criterion, self.mu, self.logvar
        real_labels = self.real_labels[:batch_size]
        for i in range(self.num_Ds):
            outputs = self.netsD[i](self.fake_imgs[i], mu)
            errG = criterion(outputs[0], real_labels)
            if len(outputs) > 1 and cfg.TRAIN.COEFF.UNCOND_LOSS > 0:
                errG_patch = cfg.TRAIN.COEFF.UNCOND_LOSS *\
                    criterion(outputs[1], real_labels)
                errG = errG + errG_patch
            errG_total = errG_total + errG
            if flag == 0:
                summary_D = summary.scalar('G_loss%d' % i, errG.item())
                self.summary_writer.add_summary(summary_D, count)

        # Compute color consistency losses
        if cfg.TRAIN.COEFF.COLOR_LOSS > 0:
            if self.num_Ds > 1:
                mu1, covariance1 = compute_mean_covariance(self.fake_imgs[-1])
                mu2, covariance2 = \
                    compute_mean_covariance(self.fake_imgs[-2].detach())
                like_mu2 = cfg.TRAIN.COEFF.COLOR_LOSS * nn.MSELoss()(mu1, mu2)
                like_cov2 = cfg.TRAIN.COEFF.COLOR_LOSS * 5 * \
                    nn.MSELoss()(covariance1, covariance2)
                errG_total = errG_total + like_mu2 + like_cov2
                if flag == 0:
                    sum_mu = summary.scalar('G_like_mu2', like_mu2.item())
                    self.summary_writer.add_summary(sum_mu, count)
                    sum_cov = summary.scalar('G_like_cov2', like_cov2.item())
                    self.summary_writer.add_summary(sum_cov, count)
            if self.num_Ds > 2:
                mu1, covariance1 = compute_mean_covariance(self.fake_imgs[-2])
                mu2, covariance2 = \
                    compute_mean_covariance(self.fake_imgs[-3].detach())
                like_mu1 = cfg.TRAIN.COEFF.COLOR_LOSS * nn.MSELoss()(mu1, mu2)
                like_cov1 = cfg.TRAIN.COEFF.COLOR_LOSS * 5 * \
                    nn.MSELoss()(covariance1, covariance2)
                errG_total = errG_total + like_mu1 + like_cov1
                if flag == 0:
                    sum_mu = summary.scalar('G_like_mu1', like_mu1.item())
                    self.summary_writer.add_summary(sum_mu, count)
                    sum_cov = summary.scalar('G_like_cov1', like_cov1.item())
                    self.summary_writer.add_summary(sum_cov, count)

        kl_loss = KL_loss(mu, logvar) * cfg.TRAIN.COEFF.KL
        errG_total = errG_total + kl_loss
        # Postpone the backward propagation
        # errG_total.backward()
        # self.optimizerG.step()
        return kl_loss, errG_total

    def train(self):
        self.netG, self.netsD, self.num_Ds,\
            self.inception_model, start_count = load_network(self.gpus)
        avg_param_G = copy_G_params(self.netG)

        self.optimizerG, self.optimizersD = \
            define_optimizers(self.netG, self.netsD)

        self.criterion = nn.BCELoss()
        self.SATcriterion = nn.CrossEntropyLoss()

        self.real_labels = Variable(torch.FloatTensor(self.batch_size).fill_(1))
        self.fake_labels = Variable(torch.FloatTensor(self.batch_size).fill_(0))

        self.gradient_one = torch.FloatTensor([1.0])
        self.gradient_half = torch.FloatTensor([0.5])

        nz = cfg.GAN.Z_DIM
        noise = Variable(torch.FloatTensor(self.batch_size, nz))
        fixed_noise = Variable(torch.FloatTensor(self.batch_size, nz).normal_(0, 1))

        # Data parameters
        data_folder = 'birds_output'  # folder with data files saved by create_input_files.py
        data_name = 'CUB_5_cap_per_img_5_min_word_freq'  # base name shared by data files

        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

        # Show, Attend, and Tell Dataloader
        train_loader = torch.utils.data.DataLoader(
                    CaptionDataset(data_folder, data_name, 'TRAIN', transform=transforms.Compose([normalize])),
                    batch_size=self.batch_size, shuffle=True, num_workers=int(cfg.WORKERS), pin_memory=True)

        if cfg.CUDA:
            self.criterion.cuda()
            self.SATcriterion.cuda() # Compute SATloss
            self.real_labels = self.real_labels.cuda()
            self.fake_labels = self.fake_labels.cuda()
            self.gradient_one = self.gradient_one.cuda()
            self.gradient_half = self.gradient_half.cuda()
            noise, fixed_noise = noise.cuda(), fixed_noise.cuda()

        predictions = []
        count = start_count
        start_epoch = start_count // (self.num_batches)
        for epoch in range(start_epoch, self.max_epoch):
            start_t = time.time()

            # for step, data in enumerate(self.data_loader, 0):
            for step, data in enumerate(zip(self.data_loader, train_loader), 0):        

                data_1 = data[0]
                _, caps, caplens = data[1]
                data = data_1

                #######################################################
                # (0) Prepare training data
                ######################################################
                self.imgs_tcpu, self.real_imgs, self.wrong_imgs, \
                    self.txt_embedding = self.prepare_data(data)

                # Testing line for real samples
                if epoch == start_epoch and step == 0:
                    print ('Checking real samples at first...')
                    save_real(self.imgs_tcpu, self.image_dir)

                #######################################################
                # (1) Generate fake images
                ######################################################
                noise.data.normal_(0, 1)
                self.fake_imgs, self.mu, self.logvar = \
                    self.netG(noise, self.txt_embedding)

                # len(self.fake_imgs) = NUM_BRANCHES
                # self.fake_imgs[0].shape = [batch_size, 3, 64, 64]
                # self.fake_imgs[1].shape = [batch_size, 3, 128, 128]
                # self.fake_imgs[2].shape = [batch_size, 3, 256, 256]

                #######################################################
                # (*) Forward fake images to SAT
                ######################################################
                from SATmodels import Encoder, DecoderWithAttention
                from torch.nn.utils.rnn import pack_padded_sequence

                fine_tune_encoder = False

                # Read word map 
                word_map_file = os.path.join(data_folder, 'WORDMAP_' + data_name + '.json') 
                with open(word_map_file, 'r') as j:
                    word_map = json.load(j)

                # Define the encoder/decoder structure for SAT model
                decoder = DecoderWithAttention(attention_dim=512,
                                               embed_dim=512,
                                               decoder_dim=512,
                                               vocab_size=len(word_map),
                                               dropout=0.5).cuda()
                decoder_optimizer = torch.optim.Adam(params=filter(lambda p: p.requires_grad, decoder.parameters()),
                                                     lr=4e-4)

                encoder = Encoder().cuda()
                encoder.fine_tune(fine_tune_encoder)
                encoder_optimizer = torch.optim.Adam(params=filter(lambda p: p.requires_grad, encoder.parameters()),
                                                     lr=1e-4) if fine_tune_encoder else None

                SATloss = 0

                # Compute the SAT loss after forwarding the SAT model
                for idx in range(len(self.fake_imgs)):
                    img = encoder(self.fake_imgs[idx])
                    scores, caps_sorted, decode_lengths, alphas, sort_ind = decoder(img, caps, caplens)
                    targets = caps_sorted[:, 1:]
                    scores, _ = pack_padded_sequence(scores, decode_lengths, batch_first=True).cuda()
                    targets, _ = pack_padded_sequence(targets, decode_lengths, batch_first=True).cuda()
                    SATloss += self.SATcriterion(scores, targets) + 1 * ((1. - alphas.sum(dim=1)) ** 2).mean()


                # Set zero_grad for encoder/decoder
                decoder_optimizer.zero_grad()
                if encoder_optimizer is not None:
                    encoder_optimizer.zero_grad()

                #######################################################
                # (2) Update D network
                ######################################################
                errD_total = 0
                for i in range(self.num_Ds):
                    errD = self.train_Dnet(i, count)
                    errD_total += errD

                #######################################################
                # (3) Update G network: maximize log(D(G(z)))
                ######################################################
                kl_loss, errG_total = self.train_Gnet(count)
                for p, avg_p in zip(self.netG.parameters(), avg_param_G):
                    avg_p.mul_(0.999).add_(0.001, p.data)

                # Combine with G and SAT first, then back propagation
                errG_total += SATloss
                errG_total.backward()
                self.optimizerG.step()

                #######################################################
                # (*) Update SAT network:
                ######################################################
                # Update weights
                decoder_optimizer.step()
                if encoder_optimizer is not None:
                    encoder_optimizer.step()

                #######################################################
                # (*) Prediction and Inception score:
                ######################################################
                pred = self.inception_model(self.fake_imgs[-1].detach())
                predictions.append(pred.data.cpu().numpy())

                if count % 100 == 0:
                    summary_D = summary.scalar('D_loss', errD_total.item())
                    summary_G = summary.scalar('G_loss', errG_total.item())
                    summary_KL = summary.scalar('KL_loss', kl_loss.item())
                    self.summary_writer.add_summary(summary_D, count)
                    self.summary_writer.add_summary(summary_G, count)
                    self.summary_writer.add_summary(summary_KL, count)

                count += 1

                #######################################################
                # (*) Save Images/Log/Model per SNAPSHOT_INTERVAL:
                ######################################################
                if count % cfg.TRAIN.SNAPSHOT_INTERVAL == 0:
                    save_model(self.netG, avg_param_G, self.netsD, count, self.model_dir)
                    # Save images
                    backup_para = copy_G_params(self.netG)
                    load_params(self.netG, avg_param_G)
                    #
                    self.fake_imgs, _, _ = self.netG(fixed_noise, self.txt_embedding)
                    save_img_results(self.imgs_tcpu, self.fake_imgs, self.num_Ds,
                                     count, self.image_dir, self.summary_writer)
                    #
                    load_params(self.netG, backup_para)

                    # Compute inception score
                    if len(predictions) > 500:
                        predictions = np.concatenate(predictions, 0)
                        mean, std = compute_inception_score(predictions, 10)
                        # print('mean:', mean, 'std', std)
                        m_incep = summary.scalar('Inception_mean', mean)
                        self.summary_writer.add_summary(m_incep, count)
                        #
                        mean_nlpp, std_nlpp = negative_log_posterior_probability(predictions, 10)
                        m_nlpp = summary.scalar('NLPP_mean', mean_nlpp)
                        self.summary_writer.add_summary(m_nlpp, count)
                        #
                        predictions = []

            end_t = time.time()
            print('''[%d/%d][%d]
                         Loss_D: %.2f Loss_G: %.2f Loss_KL: %.2f Time: %.2fs
                      '''  # D(real): %.4f D(wrong):%.4f  D(fake) %.4f
                  % (epoch, self.max_epoch, self.num_batches,
                     errD_total.item(), errG_total.item(),
                     kl_loss.item(), end_t - start_t))

        save_model(self.netG, avg_param_G, self.netsD, count, self.model_dir)
        self.summary_writer.close()

    def save_superimages(self, images_list, filenames,
                         save_dir, split_dir, imsize):
        batch_size = images_list[0].size(0)
        num_sentences = len(images_list)
        for i in range(batch_size):
            s_tmp = '%s/super/%s/%s' %\
                (save_dir, split_dir, filenames[i])
            folder = s_tmp[:s_tmp.rfind('/')]
            if not os.path.isdir(folder):
                print('Make a new folder: ', folder)
                mkdir_p(folder)
            #
            savename = '%s_%d.png' % (s_tmp, imsize)
            super_img = []
            for j in range(num_sentences):
                img = images_list[j][i]
                # print(img.size())
                img = img.view(1, 3, imsize, imsize)
                # print(img.size())
                super_img.append(img)
                # break
            super_img = torch.cat(super_img, 0)
            vutils.save_image(super_img, savename, nrow=10, normalize=True)

    def save_singleimages(self, images, filenames,
                          save_dir, split_dir, sentenceID, imsize):
        for i in range(images.size(0)):
            s_tmp = '%s/single_samples/%s/%s' %\
                (save_dir, split_dir, filenames[i])
            folder = s_tmp[:s_tmp.rfind('/')]
            if not os.path.isdir(folder):
                print('Make a new folder: ', folder)
                mkdir_p(folder)

            fullpath = '%s_%d_sentence%d.png' % (s_tmp, imsize, sentenceID)
            # range from [-1, 1] to [0, 255]
            img = images[i].add(1).div(2).mul(255).clamp(0, 255).byte()
            ndarr = img.permute(1, 2, 0).data.cpu().numpy()
            im = Image.fromarray(ndarr)
            im.save(fullpath)

    def evaluate(self, split_dir):
        if cfg.TRAIN.NET_G == '':
            print('Error: the path for morels is not found!')
        else:
            # Build and load the generator
            if split_dir == 'test':
                split_dir = 'valid'
            netG = G_NET()
            netG.apply(weights_init)
            netG = torch.nn.DataParallel(netG, device_ids=self.gpus)
            print(netG)
            # state_dict = torch.load(cfg.TRAIN.NET_G)
            state_dict = \
                torch.load(cfg.TRAIN.NET_G,
                           map_location=lambda storage, loc: storage)
            netG.load_state_dict(state_dict)
            print('Load ', cfg.TRAIN.NET_G)

            # the path to save generated images
            s_tmp = cfg.TRAIN.NET_G
            istart = s_tmp.rfind('_') + 1
            iend = s_tmp.rfind('.')
            iteration = int(s_tmp[istart:iend])
            s_tmp = s_tmp[:s_tmp.rfind('/')]
            save_dir = '%s/iteration%d' % (s_tmp, iteration)

            nz = cfg.GAN.Z_DIM
            noise = Variable(torch.FloatTensor(self.batch_size, nz))
            if cfg.CUDA:
                netG.cuda()
                noise = noise.cuda()

            # switch to evaluate mode
            netG.eval()
            for step, data in enumerate(self.data_loader, 0):
                imgs, t_embeddings, filenames = data
                if cfg.CUDA:
                    t_embeddings = Variable(t_embeddings).cuda()
                else:
                    t_embeddings = Variable(t_embeddings)
                # print(t_embeddings[:, 0, :], t_embeddings.size(1))

                embedding_dim = t_embeddings.size(1)
                batch_size = imgs[0].size(0)
                noise.data.resize_(batch_size, nz)
                noise.data.normal_(0, 1)

                fake_img_list = []
                for i in range(embedding_dim):
                    fake_imgs, _, _ = netG(noise, t_embeddings[:, i, :])
                    if cfg.TEST.B_EXAMPLE:
                        # fake_img_list.append(fake_imgs[0].data.cpu())
                        # fake_img_list.append(fake_imgs[1].data.cpu())
                        fake_img_list.append(fake_imgs[2].data.cpu())
                    else:
                        self.save_singleimages(fake_imgs[-1], filenames,
                                               save_dir, split_dir, i, 256)
                        # self.save_singleimages(fake_imgs[-2], filenames,
                        #                        save_dir, split_dir, i, 128)
                        # self.save_singleimages(fake_imgs[-3], filenames,
                        #                        save_dir, split_dir, i, 64)
                    # break
                if cfg.TEST.B_EXAMPLE:
                    # self.save_superimages(fake_img_list, filenames,
                    #                       save_dir, split_dir, 64)
                    # self.save_superimages(fake_img_list, filenames,
                    #                       save_dir, split_dir, 128)
                    self.save_superimages(fake_img_list, filenames,
                                          save_dir, split_dir, 256)

[((46, 17, 46, 54), 'torch.bmm', 'torch.bmm', ({(46, 27, 46, 34): 'img_hat', (46, 36, 46, 53): 'img_hat_transpose'}, {}), '(img_hat, img_hat_transpose)', False, 'import torch\n'), ((85, 13, 85, 30), 'six.moves.range', 'range', ({(85, 19, 85, 29): 'num_splits'}, {}), '(num_splits)', False, 'from six.moves import range\n'), ((99, 13, 99, 30), 'six.moves.range', 'range', ({(99, 19, 99, 29): 'num_splits'}, {}), '(num_splits)', False, 'from six.moves import range\n'), ((110, 11, 110, 18), 'model.G_NET', 'G_NET', ({}, {}), '()', False, 'from model import G_NET, D_NET64, D_NET128, D_NET256, D_NET512, D_NET1024, INCEPTION_V3\n'), ((112, 11, 112, 55), 'torch.nn.DataParallel', 'torch.nn.DataParallel', (), '', False, 'import torch\n'), ((151, 22, 151, 36), 'model.INCEPTION_V3', 'INCEPTION_V3', ({}, {}), '()', False, 'from model import G_NET, D_NET64, D_NET128, D_NET256, D_NET512, D_NET1024, INCEPTION_V3\n'), ((166, 13, 166, 26), 'six.moves.range', 'range', ({(166, 19, 166, 25): 'num_Ds'}, {}), '(num_Ds)', False, 'from six.moves import range\n'), ((200, 4, 202, 23), 'torchvision.utils.save_image', 'vutils.save_image', (), '', True, 'import torchvision.utils as vutils\n'), ((204, 19, 204, 56), 'numpy.transpose', 'np.transpose', ({(204, 32, 204, 44): 'real_img_set', (204, 46, 204, 55): '(1, 2, 0)'}, {}), '(real_img_set, (1, 2, 0))', True, 'import numpy as np\n'), ((207, 19, 207, 58), 'tensorboard.summary.image', 'summary.image', ({(207, 33, 207, 43): '"""real_img"""', (207, 45, 207, 57): 'real_img_set'}, {}), "('real_img', real_img_set)", False, 'from tensorboard import summary\n'), ((217, 4, 219, 23), 'torchvision.utils.save_image', 'vutils.save_image', (), '', True, 'import torchvision.utils as vutils\n'), ((221, 19, 221, 56), 'numpy.transpose', 'np.transpose', ({(221, 32, 221, 44): 'real_img_set', (221, 46, 221, 55): '(1, 2, 0)'}, {}), '(real_img_set, (1, 2, 0))', True, 'import numpy as np\n'), ((224, 19, 224, 58), 'tensorboard.summary.image', 'summary.image', ({(224, 33, 224, 43): '"""real_img"""', (224, 45, 224, 57): 'real_img_set'}, {}), "('real_img', real_img_set)", False, 'from tensorboard import summary\n'), ((227, 13, 227, 28), 'six.moves.range', 'range', ({(227, 19, 227, 27): 'num_imgs'}, {}), '(num_imgs)', False, 'from six.moves import range\n'), ((62, 8, 62, 46), 'torch.nn.init.orthogonal', 'nn.init.orthogonal', ({(62, 27, 62, 40): 'm.weight.data', (62, 42, 62, 45): '(1.0)'}, {}), '(m.weight.data, 1.0)', True, 'import torch.nn as nn\n'), ((93, 11, 93, 26), 'numpy.mean', 'np.mean', ({(93, 19, 93, 25): 'scores'}, {}), '(scores)', True, 'import numpy as np\n'), ((93, 28, 93, 42), 'numpy.std', 'np.std', ({(93, 35, 93, 41): 'scores'}, {}), '(scores)', True, 'import numpy as np\n'), ((104, 17, 104, 32), 'numpy.mean', 'np.mean', ({(104, 25, 104, 31): 'result'}, {}), '(result)', True, 'import numpy as np\n'), ((106, 11, 106, 26), 'numpy.mean', 'np.mean', ({(106, 19, 106, 25): 'scores'}, {}), '(scores)', True, 'import numpy as np\n'), ((106, 28, 106, 42), 'numpy.std', 'np.std', ({(106, 35, 106, 41): 'scores'}, {}), '(scores)', True, 'import numpy as np\n'), ((130, 19, 130, 67), 'torch.nn.DataParallel', 'torch.nn.DataParallel', (), '', False, 'import torch\n'), ((136, 21, 136, 48), 'torch.load', 'torch.load', ({(136, 32, 136, 47): 'cfg.TRAIN.NET_G'}, {}), '(cfg.TRAIN.NET_G)', False, 'import torch\n'), ((141, 15, 141, 41), 'miscc.config.cfg.TRAIN.NET_G.rfind', 'cfg.TRAIN.NET_G.rfind', ({(141, 37, 141, 40): '"""."""'}, {}), "('.')", False, 'from miscc.config import cfg\n'), ((231, 8, 233, 50), 'torchvision.utils.save_image', 'vutils.save_image', (), '', True, 'import torchvision.utils as vutils\n'), ((237, 23, 237, 60), 'numpy.transpose', 'np.transpose', ({(237, 36, 237, 48): 'fake_img_set', (237, 50, 237, 59): '(1, 2, 0)'}, {}), '(fake_img_set, (1, 2, 0))', True, 'import numpy as np\n'), ((241, 23, 241, 68), 'tensorboard.summary.image', 'summary.image', ({(241, 37, 241, 53): "'fake_img%d' % i", (241, 55, 241, 67): 'fake_img_set'}, {}), "('fake_img%d' % i, fake_img_set)", False, 'from tensorboard import summary\n'), ((258, 17, 258, 38), 'miscc.config.cfg.GPU_ID.split', 'cfg.GPU_ID.split', ({(258, 34, 258, 37): '""","""'}, {}), "(',')", False, 'from miscc.config import cfg\n'), ((261, 8, 261, 43), 'torch.cuda.set_device', 'torch.cuda.set_device', ({(261, 30, 261, 42): 'self.gpus[0]'}, {}), '(self.gpus[0])', False, 'import torch\n'), ((279, 17, 279, 35), 'six.moves.range', 'range', ({(279, 23, 279, 34): 'self.num_Ds'}, {}), '(self.num_Ds)', False, 'from six.moves import range\n'), ((342, 17, 342, 35), 'six.moves.range', 'range', ({(342, 23, 342, 34): 'self.num_Ds'}, {}), '(self.num_Ds)', False, 'from six.moves import range\n'), ((398, 25, 398, 37), 'torch.nn.BCELoss', 'nn.BCELoss', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((399, 28, 399, 49), 'torch.nn.CrossEntropyLoss', 'nn.CrossEntropyLoss', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((404, 28, 404, 52), 'torch.FloatTensor', 'torch.FloatTensor', ({(404, 46, 404, 51): '[1.0]'}, {}), '([1.0])', False, 'import torch\n'), ((405, 29, 405, 53), 'torch.FloatTensor', 'torch.FloatTensor', ({(405, 47, 405, 52): '[0.5]'}, {}), '([0.5])', False, 'import torch\n'), ((415, 20, 415, 95), 'torchvision.transforms.Normalize', 'transforms.Normalize', (), '', True, 'import torchvision.transforms as transforms\n'), ((434, 21, 434, 55), 'six.moves.range', 'range', ({(434, 27, 434, 38): 'start_epoch', (434, 40, 434, 54): 'self.max_epoch'}, {}), '(start_epoch, self.max_epoch)', False, 'from six.moves import range\n'), ((599, 17, 599, 34), 'six.moves.range', 'range', ({(599, 23, 599, 33): 'batch_size'}, {}), '(batch_size)', False, 'from six.moves import range\n'), ((55, 10, 55, 33), 'torch.mean', 'torch.mean', ({(55, 21, 55, 32): 'KLD_element'}, {}), '(KLD_element)', False, 'import torch\n'), ((91, 21, 91, 34), 'numpy.sum', 'np.sum', ({(91, 28, 91, 30): 'kl', (91, 32, 91, 33): '1'}, {}), '(kl, 1)', True, 'import numpy as np\n'), ((92, 22, 92, 32), 'numpy.exp', 'np.exp', ({(92, 29, 92, 31): 'kl'}, {}), '(kl)', True, 'import numpy as np\n'), ((117, 21, 117, 30), 'model.D_NET64', 'D_NET64', ({}, {}), '()', False, 'from model import G_NET, D_NET64, D_NET128, D_NET256, D_NET512, D_NET1024, INCEPTION_V3\n'), ((119, 21, 119, 31), 'model.D_NET128', 'D_NET128', ({}, {}), '()', False, 'from model import G_NET, D_NET64, D_NET128, D_NET256, D_NET512, D_NET1024, INCEPTION_V3\n'), ((121, 21, 121, 31), 'model.D_NET256', 'D_NET256', ({}, {}), '()', False, 'from model import G_NET, D_NET64, D_NET128, D_NET256, D_NET512, D_NET1024, INCEPTION_V3\n'), ((123, 21, 123, 31), 'model.D_NET512', 'D_NET512', ({}, {}), '()', False, 'from model import G_NET, D_NET64, D_NET128, D_NET256, D_NET512, D_NET1024, INCEPTION_V3\n'), ((125, 21, 125, 32), 'model.D_NET1024', 'D_NET1024', ({}, {}), '()', False, 'from model import G_NET, D_NET64, D_NET128, D_NET256, D_NET512, D_NET1024, INCEPTION_V3\n'), ((140, 17, 140, 43), 'miscc.config.cfg.TRAIN.NET_G.rfind', 'cfg.TRAIN.NET_G.rfind', ({(140, 39, 140, 42): '"""_"""'}, {}), "('_')", False, 'from miscc.config import cfg\n'), ((148, 25, 148, 70), 'torch.load', 'torch.load', ({(148, 36, 148, 69): "'%s%d.pth' % (cfg.TRAIN.NET_D, i)"}, {}), "('%s%d.pth' % (cfg.TRAIN.NET_D, i))", False, 'import torch\n'), ((203, 19, 203, 45), 'torchvision.utils.make_grid', 'vutils.make_grid', ({(203, 36, 203, 44): 'real_img'}, {}), '(real_img)', True, 'import torchvision.utils as vutils\n'), ((220, 19, 220, 45), 'torchvision.utils.make_grid', 'vutils.make_grid', ({(220, 36, 220, 44): 'real_img'}, {}), '(real_img)', True, 'import torchvision.utils as vutils\n'), ((250, 29, 250, 62), 'os.path.join', 'os.path.join', ({(250, 42, 250, 52): 'output_dir', (250, 54, 250, 61): '"""Model"""'}, {}), "(output_dir, 'Model')", False, 'import os\n'), ((251, 29, 251, 62), 'os.path.join', 'os.path.join', ({(251, 42, 251, 52): 'output_dir', (251, 54, 251, 61): '"""Image"""'}, {}), "(output_dir, 'Image')", False, 'import os\n'), ((252, 27, 252, 58), 'os.path.join', 'os.path.join', ({(252, 40, 252, 50): 'output_dir', (252, 52, 252, 57): '"""Log"""'}, {}), "(output_dir, 'Log')", False, 'import os\n'), ((253, 12, 253, 35), 'miscc.utils.mkdir_p', 'mkdir_p', ({(253, 20, 253, 34): 'self.model_dir'}, {}), '(self.model_dir)', False, 'from miscc.utils import mkdir_p\n'), ((254, 12, 254, 35), 'miscc.utils.mkdir_p', 'mkdir_p', ({(254, 20, 254, 34): 'self.image_dir'}, {}), '(self.image_dir)', False, 'from miscc.utils import mkdir_p\n'), ((255, 12, 255, 33), 'miscc.utils.mkdir_p', 'mkdir_p', ({(255, 20, 255, 32): 'self.log_dir'}, {}), '(self.log_dir)', False, 'from miscc.utils import mkdir_p\n'), ((256, 34, 256, 58), 'tensorboard.FileWriter', 'FileWriter', ({(256, 45, 256, 57): 'self.log_dir'}, {}), '(self.log_dir)', False, 'from tensorboard import FileWriter\n'), ((278, 25, 278, 46), 'torch.autograd.Variable', 'Variable', ({(278, 34, 278, 45): 't_embedding'}, {}), '(t_embedding)', False, 'from torch.autograd import Variable\n'), ((408, 25, 408, 63), 'torch.FloatTensor', 'torch.FloatTensor', ({(408, 43, 408, 58): 'self.batch_size', (408, 60, 408, 62): 'nz'}, {}), '(self.batch_size, nz)', False, 'import torch\n'), ((435, 22, 435, 33), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((584, 20, 584, 31), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((609, 21, 609, 41), 'six.moves.range', 'range', ({(609, 27, 609, 40): 'num_sentences'}, {}), '(num_sentences)', False, 'from six.moves import range\n'), ((616, 24, 616, 47), 'torch.cat', 'torch.cat', ({(616, 34, 616, 43): 'super_img', (616, 45, 616, 46): '0'}, {}), '(super_img, 0)', False, 'import torch\n'), ((617, 12, 617, 75), 'torchvision.utils.save_image', 'vutils.save_image', (), '', True, 'import torchvision.utils as vutils\n'), ((633, 17, 633, 39), 'PIL.Image.fromarray', 'Image.fromarray', ({(633, 33, 633, 38): 'ndarr'}, {}), '(ndarr)', False, 'from PIL import Image, ImageFont, ImageDraw\n'), ((643, 19, 643, 26), 'model.G_NET', 'G_NET', ({}, {}), '()', False, 'from model import G_NET, D_NET64, D_NET128, D_NET256, D_NET512, D_NET1024, INCEPTION_V3\n'), ((645, 19, 645, 68), 'torch.nn.DataParallel', 'torch.nn.DataParallel', (), '', False, 'import torch\n'), ((649, 16, 650, 69), 'torch.load', 'torch.load', (), '', False, 'import torch\n'), ((67, 8, 67, 46), 'torch.nn.init.orthogonal', 'nn.init.orthogonal', ({(67, 27, 67, 40): 'm.weight.data', (67, 42, 67, 45): '(1.0)'}, {}), '(m.weight.data, 1.0)', True, 'import torch.nn as nn\n'), ((90, 13, 90, 25), 'numpy.log', 'np.log', ({(90, 20, 90, 24): 'part'}, {}), '(part)', True, 'import numpy as np\n'), ((103, 30, 103, 45), 'numpy.max', 'np.max', ({(103, 37, 103, 41): 'part', (103, 43, 103, 44): '(1)'}, {}), '(part, 1)', True, 'import numpy as np\n'), ((476, 32, 476, 91), 'os.path.join', 'os.path.join', ({(476, 45, 476, 56): 'data_folder', (476, 58, 476, 90): "'WORDMAP_' + data_name + '.json'"}, {}), "(data_folder, 'WORDMAP_' + data_name + '.json')", False, 'import os\n'), ((515, 25, 515, 43), 'six.moves.range', 'range', ({(515, 31, 515, 42): 'self.num_Ds'}, {}), '(self.num_Ds)', False, 'from six.moves import range\n'), ((603, 19, 603, 40), 'os.path.isdir', 'os.path.isdir', ({(603, 33, 603, 39): 'folder'}, {}), '(folder)', False, 'import os\n'), ((605, 16, 605, 31), 'miscc.utils.mkdir_p', 'mkdir_p', ({(605, 24, 605, 30): 'folder'}, {}), '(folder)', False, 'from miscc.utils import mkdir_p\n'), ((625, 19, 625, 40), 'os.path.isdir', 'os.path.isdir', ({(625, 33, 625, 39): 'folder'}, {}), '(folder)', False, 'import os\n'), ((627, 16, 627, 31), 'miscc.utils.mkdir_p', 'mkdir_p', ({(627, 24, 627, 30): 'folder'}, {}), '(folder)', False, 'from miscc.utils import mkdir_p\n'), ((663, 29, 663, 67), 'torch.FloatTensor', 'torch.FloatTensor', ({(663, 47, 663, 62): 'self.batch_size', (663, 64, 663, 66): 'nz'}, {}), '(self.batch_size, nz)', False, 'import torch\n'), ((684, 25, 684, 45), 'six.moves.range', 'range', ({(684, 31, 684, 44): 'embedding_dim'}, {}), '(embedding_dim)', False, 'from six.moves import range\n'), ((276, 25, 276, 46), 'torch.autograd.Variable', 'Variable', ({(276, 34, 276, 45): 't_embedding'}, {}), '(t_embedding)', False, 'from torch.autograd import Variable\n'), ((284, 34, 284, 51), 'torch.autograd.Variable', 'Variable', ({(284, 43, 284, 50): 'imgs[i]'}, {}), '(imgs[i])', False, 'from torch.autograd import Variable\n'), ((285, 35, 285, 54), 'torch.autograd.Variable', 'Variable', ({(285, 44, 285, 53): 'w_imgs[i]'}, {}), '(w_imgs[i])', False, 'from torch.autograd import Variable\n'), ((401, 36, 401, 70), 'torch.FloatTensor', 'torch.FloatTensor', ({(401, 54, 401, 69): 'self.batch_size'}, {}), '(self.batch_size)', False, 'import torch\n'), ((402, 36, 402, 70), 'torch.FloatTensor', 'torch.FloatTensor', ({(402, 54, 402, 69): 'self.batch_size'}, {}), '(self.batch_size)', False, 'import torch\n'), ((409, 31, 409, 69), 'torch.FloatTensor', 'torch.FloatTensor', ({(409, 49, 409, 64): 'self.batch_size', (409, 66, 409, 68): 'nz'}, {}), '(self.batch_size, nz)', False, 'import torch\n'), ((419, 78, 419, 109), 'torchvision.transforms.Compose', 'transforms.Compose', ({(419, 97, 419, 108): '[normalize]'}, {}), '([normalize])', True, 'import torchvision.transforms as transforms\n'), ((675, 35, 675, 57), 'torch.autograd.Variable', 'Variable', ({(675, 44, 675, 56): 't_embeddings'}, {}), '(t_embeddings)', False, 'from torch.autograd import Variable\n'), ((90, 50, 90, 66), 'numpy.mean', 'np.mean', ({(90, 58, 90, 62): 'part', (90, 64, 90, 65): '(0)'}, {}), '(part, 0)', True, 'import numpy as np\n'), ((235, 23, 235, 54), 'torchvision.utils.make_grid', 'vutils.make_grid', ({(235, 40, 235, 53): 'fake_img.data'}, {}), '(fake_img.data)', True, 'import torchvision.utils as vutils\n'), ((360, 56, 360, 68), 'torch.nn.MSELoss', 'nn.MSELoss', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((362, 20, 362, 32), 'torch.nn.MSELoss', 'nn.MSELoss', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((373, 56, 373, 68), 'torch.nn.MSELoss', 'nn.MSELoss', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((375, 20, 375, 32), 'torch.nn.MSELoss', 'nn.MSELoss', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((489, 26, 489, 35), 'SATmodels.Encoder', 'Encoder', ({}, {}), '()', False, 'from SATmodels import Encoder, DecoderWithAttention\n'), ((572, 38, 572, 68), 'numpy.concatenate', 'np.concatenate', ({(572, 53, 572, 64): 'predictions', (572, 66, 572, 67): '0'}, {}), '(predictions, 0)', True, 'import numpy as np\n'), ((575, 34, 575, 72), 'tensorboard.summary.scalar', 'summary.scalar', ({(575, 49, 575, 65): '"""Inception_mean"""', (575, 67, 575, 71): 'mean'}, {}), "('Inception_mean', mean)", False, 'from tensorboard import summary\n'), ((579, 33, 579, 71), 'tensorboard.summary.scalar', 'summary.scalar', ({(579, 48, 579, 59): '"""NLPP_mean"""', (579, 61, 579, 70): 'mean_nlpp'}, {}), "('NLPP_mean', mean_nlpp)", False, 'from tensorboard import summary\n'), ((281, 34, 281, 51), 'torch.autograd.Variable', 'Variable', ({(281, 43, 281, 50): 'imgs[i]'}, {}), '(imgs[i])', False, 'from torch.autograd import Variable\n'), ((282, 35, 282, 54), 'torch.autograd.Variable', 'Variable', ({(282, 44, 282, 53): 'w_imgs[i]'}, {}), '(w_imgs[i])', False, 'from torch.autograd import Variable\n'), ((501, 32, 501, 94), 'torch.nn.utils.rnn.pack_padded_sequence', 'pack_padded_sequence', (), '', False, 'from torch.nn.utils.rnn import pack_padded_sequence\n'), ((502, 33, 502, 96), 'torch.nn.utils.rnn.pack_padded_sequence', 'pack_padded_sequence', (), '', False, 'from torch.nn.utils.rnn import pack_padded_sequence\n'), ((673, 35, 673, 57), 'torch.autograd.Variable', 'Variable', ({(673, 44, 673, 56): 't_embeddings'}, {}), '(t_embeddings)', False, 'from torch.autograd import Variable\n')]