# PyTorch StudioGAN: https://github.com/POSTECH-CVLab/PyTorch-StudioGAN
# The MIT License (MIT)
# See license file or visit https://github.com/POSTECH-CVLab/PyTorch-StudioGAN for details

# train_eval.py

import numpy as np
import sys
import glob
from scipy import ndimage
from os.path import join
from PIL import Image
from tqdm import tqdm
from datetime import datetime

from metrics.IS import calculate_incep_score
from metrics.FID import calculate_fid_score
from metrics.F_beta import calculate_f_beta_score
from metrics.Accuracy import calculate_accuracy
from utils.ada import augment
from utils.biggan_utils import interp
from utils.sample import sample_latents, sample_1hot, make_mask, target_class_sampler
from utils.misc import *
from utils.losses import calc_derv4gp, calc_derv4dra, calc_derv, latent_optimise
from utils.losses import Conditional_Contrastive_loss, Proxy_NCA_loss, NT_Xent_loss
from utils.diff_aug import DiffAugment
from utils.cr_diff_aug import CR_DiffAug

import torch
import torch.nn as nn
from torch.nn import DataParallel
import torch.nn.functional as F
import torchvision
from torchvision import transforms

SAVE_FORMAT = 'step={step:0>3}-Inception_mean={Inception_mean:<.4}-Inception_std={Inception_std:<.4}-FID={FID:<.5}.pth'

LOG_FORMAT = (
    "Step: {step:>7} "
    "Progress: {progress:<.1%} "
    "Elapsed: {elapsed} "
    "temperature: {temperature:<.6} "
    "ada_p: {ada_p:<.6} "
    "Discriminator_loss: {dis_loss:<.6} "
    "Generator_loss: {gen_loss:<.6} "
)

def set_temperature(conditional_strategy, tempering_type, start_temperature, end_temperature, step_count, tempering_step, total_step):
    if conditional_strategy == 'ContraGAN':
        if tempering_type == 'continuous':
            t = start_temperature + step_count*(end_temperature - start_temperature)/total_step
        elif tempering_type == 'discrete':
            tempering_interval = total_step//(tempering_step + 1)
            t = start_temperature + \
                (step_count//tempering_interval)*(end_temperature-start_temperature)/tempering_step
        else:
            t = start_temperature
    else:
        t = 'no'
    return t

class Train_Eval(object):
    def __init__(self, run_name, best_step, dataset_name, eval_type, logger, writer, n_gpus, gen_model, dis_model, inception_model,
                 Gen_copy, Gen_ema, train_dataset, eval_dataset, train_dataloader, eval_dataloader, freeze_layers, conditional_strategy,
                 pos_collected_numerator, z_dim, num_classes, hypersphere_dim, d_spectral_norm, g_spectral_norm, G_optimizer, D_optimizer,
                 batch_size, g_steps_per_iter, d_steps_per_iter, accumulation_steps, total_step, G_loss, D_loss, contrastive_lambda, margin,
                 tempering_type, tempering_step, start_temperature, end_temperature, weight_clipping_for_dis, weight_clipping_bound,
                 gradient_penalty_for_dis, gradient_penalty_lambda, deep_regret_analysis_for_dis, regret_penalty_lambda, cr, cr_lambda, bcr,
                 real_lambda, fake_lambda, zcr, gen_lambda, dis_lambda, sigma_noise, diff_aug, ada, prev_ada_p, ada_target, ada_length, prior,
                 truncated_factor, ema, latent_op, latent_op_rate, latent_op_step, latent_op_step4eval, latent_op_alpha, latent_op_beta,
                 latent_norm_reg_weight, default_device, print_every, save_every, checkpoint_dir, evaluate, mu, sigma, best_fid,
                 best_fid_checkpoint_path, mixed_precision, train_config, model_config, gamma, steps):

        self.run_name = run_name
        self.best_step = best_step
        self.dataset_name = dataset_name
        self.eval_type = eval_type
        self.logger = logger
        self.writer = writer
        self.n_gpus = n_gpus

        self.gen_model = gen_model
        self.dis_model = dis_model
        self.inception_model = inception_model
        self.Gen_copy = Gen_copy
        self.Gen_ema = Gen_ema

        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
        self.train_dataloader = train_dataloader
        self.eval_dataloader = eval_dataloader

        self.freeze_layers = freeze_layers

        self.conditional_strategy = conditional_strategy
        self.pos_collected_numerator = pos_collected_numerator
        self.z_dim = z_dim
        self.num_classes = num_classes
        self.hypersphere_dim = hypersphere_dim
        self.d_spectral_norm = d_spectral_norm
        self.g_spectral_norm = g_spectral_norm

        self.G_optimizer = G_optimizer
        self.D_optimizer = D_optimizer
        self.batch_size = batch_size
        self.g_steps_per_iter = g_steps_per_iter
        self.d_steps_per_iter = d_steps_per_iter
        self.accumulation_steps = accumulation_steps
        self.total_step = total_step

        self.G_loss = G_loss
        self.D_loss = D_loss
        self.contrastive_lambda = contrastive_lambda
        self.margin = margin
        self.tempering_type = tempering_type
        self.tempering_step = tempering_step
        self.start_temperature = start_temperature
        self.end_temperature = end_temperature
        self.weight_clipping_for_dis = weight_clipping_for_dis
        self.weight_clipping_bound = weight_clipping_bound
        self.gradient_penalty_for_dis = gradient_penalty_for_dis
        self.gradient_penalty_lambda = gradient_penalty_lambda
        self.deep_regret_analysis_for_dis = deep_regret_analysis_for_dis
        self.regret_penalty_lambda = regret_penalty_lambda
        self.cr = cr
        self.cr_lambda = cr_lambda
        self.bcr = bcr
        self.real_lambda = real_lambda
        self.fake_lambda = fake_lambda
        self.zcr = zcr
        self.gen_lambda = gen_lambda
        self.dis_lambda = dis_lambda
        self.sigma_noise = sigma_noise

        self.diff_aug = diff_aug
        self.ada = ada
        self.prev_ada_p = prev_ada_p
        self.ada_target = ada_target
        self.ada_length = ada_length
        self.prior = prior
        self.truncated_factor = truncated_factor
        self.ema = ema
        self.latent_op = latent_op
        self.latent_op_rate = latent_op_rate
        self.latent_op_step = latent_op_step
        self.latent_op_step4eval = latent_op_step4eval
        self.latent_op_alpha = latent_op_alpha
        self.latent_op_beta = latent_op_beta
        self.latent_norm_reg_weight = latent_norm_reg_weight

        self.default_device = default_device
        self.print_every = print_every
        self.save_every = save_every
        self.checkpoint_dir = checkpoint_dir
        self.evaluate = evaluate
        self.mu = mu
        self.sigma = sigma
        self.best_fid = best_fid
        self.best_fid_checkpoint_path = best_fid_checkpoint_path
        self.mixed_precision = mixed_precision
        self.train_config = train_config
        self.model_config = model_config

        self.start_time = datetime.now()
        self.l2_loss = torch.nn.MSELoss()
        self.ce_loss = torch.nn.CrossEntropyLoss()
        self.policy = "color,translation,cutout"

        self.steps = steps
        self.gamma = gamma

        sampler = define_sampler(self.dataset_name, self.conditional_strategy)

        check_flag_1(self.tempering_type, self.pos_collected_numerator, self.conditional_strategy, self.diff_aug, self.ada,
                     self.mixed_precision, self.gradient_penalty_for_dis, self.deep_regret_analysis_for_dis, self.cr, self.bcr, self.zcr)

        if self.conditional_strategy == 'ContraGAN':
            self.contrastive_criterion = Conditional_Contrastive_loss(self.default_device, self.batch_size, self.pos_collected_numerator)

        elif self.conditional_strategy == 'Proxy_NCA_GAN':
            if isinstance(self.dis_model, DataParallel):
                self.embedding_layer = self.dis_model.module.embedding
            else:
                self.embedding_layer = self.dis_model.embedding
            self.NCA_criterion = Proxy_NCA_loss(self.default_device, self.embedding_layer, self.num_classes, self.batch_size)

        elif self.conditional_strategy == 'NT_Xent_GAN':
            self.NT_Xent_criterion = NT_Xent_loss(self.default_device, self.batch_size)
        else:
            pass

        if self.mixed_precision:
            self.scaler = torch.cuda.amp.GradScaler()

        if self.dataset_name in ["imagenet"]:
            self.num_eval = {'train':50000, 'valid':50000}
        elif self.dataset_name in ["imagenet_less_0.25"]:
            self.num_eval = {'train':50000, 'valid':50000}
        elif self.dataset_name in ["imagenet_less"]:
            self.num_eval = {'train':50000, 'valid':50000}
        elif self.dataset_name == "tiny_imagenet":
            self.num_eval = {'train':50000, 'valid':10000}
        elif self.dataset_name == "cifar10":
            self.num_eval = {'train':50000, 'test':10000}
        elif self.dataset_name == "cifar10_less":
            self.num_eval = {'train':len(self.train_dataset.data), 'valid':len(self.eval_dataset.data), 'test':len(self.eval_dataset.data)}
        elif self.dataset_name in ["cifar100_less"]:
            self.num_eval = {'train':len(self.train_dataset.data), 'valid':len(self.eval_dataset.data), 'test':len(self.eval_dataset.data)}
        elif self.dataset_name == "custom":
            num_train_images = len(self.train_dataset.data)
            num_eval_images = len(self.eval_dataset.data)
            self.num_eval = {'train':num_train_images, 'valid':num_eval_images}
        else:
            raise NotImplementedError

    ################################################################################################################################
    def train(self, current_step, total_step):
        self.dis_model.train()
        self.gen_model.train()
        if self.Gen_copy is not None:
            self.Gen_copy.train()

        self.logger.info('Start training....')
        step_count = current_step
        train_iter = iter(self.train_dataloader)

        if self.ada:
            self.ada_augment = torch.tensor([0.0, 0.0], device = self.default_device)
            if self.prev_ada_p is not None:
                self.ada_aug_p = self.prev_ada_p
            else:
                self.ada_aug_p = 0.0
            self.ada_aug_step = self.ada_target/self.ada_length
        else:
            self.ada_aug_p = 'No'

        while step_count <= total_step:
            # ================== TRAIN D ================== #
            toggle_grad(self.dis_model, True, freeze_layers=self.freeze_layers)
            toggle_grad(self.gen_model, False, freeze_layers=-1)
            t = set_temperature(self.conditional_strategy, self.tempering_type, self.start_temperature, self.end_temperature, step_count, self.tempering_step, total_step)
            for step_index in range(self.d_steps_per_iter):
                self.D_optimizer.zero_grad()
                for acml_index in range(self.accumulation_steps):
                    try:
                        real_images, real_labels = next(train_iter)
                    except StopIteration:
                        train_iter = iter(self.train_dataloader)
                        real_images, real_labels = next(train_iter)

                    real_images, real_labels = real_images.to(self.default_device), real_labels.to(self.default_device)
                    with torch.cuda.amp.autocast() if self.mixed_precision else dummy_context_mgr() as mpc:
                        if self.diff_aug:
                            real_images = DiffAugment(real_images, policy=self.policy)
                        if self.ada:
                            real_images, _ = augment(real_images, self.ada_aug_p)

                        if self.zcr:
                            zs, fake_labels, zs_t = sample_latents(self.prior, self.batch_size, self.z_dim, 1, self.num_classes,
                                                                   self.sigma_noise, self.default_device)
                        else:
                            zs, fake_labels = sample_latents(self.prior, self.batch_size, self.z_dim, 1, self.num_classes,
                                                             None, self.default_device)
                        if self.latent_op:
                            zs = latent_optimise(zs, fake_labels, self.gen_model, self.dis_model, self.conditional_strategy,
                                                 self.latent_op_step, self.latent_op_rate, self.latent_op_alpha, self.latent_op_beta,
                                                 False, self.default_device)

                        fake_images = self.gen_model(zs, fake_labels)
                        if self.diff_aug:
                            fake_images = DiffAugment(fake_images, policy=self.policy)
                        if self.ada:
                            fake_images, _ = augment(fake_images, self.ada_aug_p)

                        if self.conditional_strategy == "ACGAN":
                            cls_out_real, dis_out_real = self.dis_model(real_images, real_labels)
                            cls_out_fake, dis_out_fake = self.dis_model(fake_images, fake_labels)
                        elif self.conditional_strategy == "ProjGAN" or self.conditional_strategy == "no":
                            dis_out_real = self.dis_model(real_images, real_labels)
                            dis_out_fake = self.dis_model(fake_images, fake_labels)
                        elif self.conditional_strategy in ["NT_Xent_GAN", "Proxy_NCA_GAN", "ContraGAN"]:
                            real_cls_mask = make_mask(real_labels, self.num_classes, self.default_device)
                            cls_proxies_real, cls_embed_real, dis_out_real = self.dis_model(real_images, real_labels)
                            cls_proxies_fake, cls_embed_fake, dis_out_fake = self.dis_model(fake_images, fake_labels)
                        elif self.conditional_strategy == 'ProjGAN_adv':
                            dis_out_real_prefc = self.dis_model(real_images, real_labels, fc=False)
                            dis_out_fake_prefc = self.dis_model(fake_images, fake_labels, fc=False)

                            loss_real = lambda x: torch.mean(F.relu(1. - x))
                            loss_fake = lambda x: torch.mean(F.relu(1. + x))
                            dis_out_real_prefc_adv = PGD(dis_out_real_prefc, real_labels, loss_real, self.dis_model, steps=self.steps, gamma=self.gamma)
                            dis_out_fake_prefc_adv = PGD(dis_out_fake_prefc, fake_labels, loss_real, self.dis_model, steps=self.steps, gamma=self.gamma)

                            fake_images = fake_images.detach()
                            dis_out_real_prefc = self.dis_model(real_images, real_labels, fc=False, only_fc=False)
                            dis_out_fake_prefc = self.dis_model(fake_images, fake_labels, fc=False, only_fc=False)

                            dis_out_real = self.dis_model(dis_out_real_prefc, real_labels, only_fc=True, fc=True)
                            dis_out_fake = self.dis_model(dis_out_fake_prefc, fake_labels, only_fc=True, fc=True)

                            dis_out_real_adv = self.dis_model(dis_out_real_prefc_adv, real_labels, only_fc=True)
                            dis_out_fake_adv = self.dis_model(dis_out_fake_prefc_adv, fake_labels, only_fc=True)

                        else:
                            raise NotImplementedError

                        #if self.conditional_strategy != 'ProjGAN_adv':
                        if self.conditional_strategy != 'ProjGAN_adv':
                            dis_acml_loss = self.D_loss(dis_out_real, dis_out_fake)
                        else:
                            dis_acml_loss = (self.D_loss(dis_out_real, dis_out_fake) + self.D_loss(dis_out_real_adv, dis_out_fake_adv)) / 2

                        if self.conditional_strategy == "ACGAN":
                            dis_acml_loss += (self.ce_loss(cls_out_real, real_labels) + self.ce_loss(cls_out_fake, fake_labels))
                        elif self.conditional_strategy == "NT_Xent_GAN":
                            real_images_aug = CR_DiffAug(real_images)
                            _, cls_embed_real_aug, dis_out_real_aug = self.dis_model(real_images_aug, real_labels)
                            dis_acml_loss += self.contrastive_lambda*self.NT_Xent_criterion(cls_embed_real, cls_embed_real_aug, t)
                        elif self.conditional_strategy == "Proxy_NCA_GAN":
                            dis_acml_loss += self.contrastive_lambda*self.NCA_criterion(cls_embed_real, cls_proxies_real, real_labels)
                        elif self.conditional_strategy == "ContraGAN":
                            dis_acml_loss += self.contrastive_lambda*self.contrastive_criterion(cls_embed_real, cls_proxies_real,
                                                                                real_cls_mask, real_labels, t, self.margin)
                        else:
                            pass

                        if self.cr:
                            real_images_aug = CR_DiffAug(real_images)
                            if self.conditional_strategy == "ACGAN":
                                cls_out_real_aug, dis_out_real_aug = self.dis_model(real_images_aug, real_labels)
                                cls_consistency_loss = self.l2_loss(cls_out_real, cls_out_real_aug)
                            elif self.conditional_strategy == "ProjGAN" or self.conditional_strategy == "no":
                                dis_out_real_aug = self.dis_model(real_images_aug, real_labels)
                            elif self.conditional_strategy in ["NT_Xent_GAN", "Proxy_NCA_GAN", "ContraGAN"]:
                                _, cls_embed_real_aug, dis_out_real_aug = self.dis_model(real_images_aug, real_labels)
                                cls_consistency_loss = self.l2_loss(cls_embed_real, cls_embed_real_aug)
                            elif self.conditional_strategy == "ProjGAN_adv":
                                dis_out_real_aug = self.dis_model(real_images_aug, real_labels)
                            else:
                                raise NotImplementedError

                            consistency_loss = self.l2_loss(dis_out_real, dis_out_real_aug)
                            if self.conditional_strategy in ["ACGAN", "NT_Xent_GAN", "Proxy_NCA_GAN", "ContraGAN"]:
                                consistency_loss += cls_consistency_loss
                            dis_acml_loss += self.cr_lambda*consistency_loss

                        if self.bcr:
                            real_images_aug = CR_DiffAug(real_images)
                            fake_images_aug = CR_DiffAug(fake_images)
                            if self.conditional_strategy == "ACGAN":
                                cls_out_real_aug, dis_out_real_aug = self.dis_model(real_images_aug, real_labels)
                                cls_out_fake_aug, dis_out_fake_aug = self.dis_model(fake_images_aug, fake_labels)
                                cls_bcr_real_loss = self.l2_loss(cls_out_real, cls_out_real_aug)
                                cls_bcr_fake_loss = self.l2_loss(cls_out_fake, cls_out_fake_aug)
                            elif self.conditional_strategy == "ProjGAN" or self.conditional_strategy == "no":
                                dis_out_real_aug = self.dis_model(real_images_aug, real_labels)
                                dis_out_fake_aug = self.dis_model(fake_images_aug, fake_labels)
                            elif self.conditional_strategy in ["ContraGAN", "Proxy_NCA_GAN", "NT_Xent_GAN"]:
                                cls_proxies_real_aug, cls_embed_real_aug, dis_out_real_aug = self.dis_model(real_images_aug, real_labels)
                                cls_proxies_fake_aug, cls_embed_fake_aug, dis_out_fake_aug = self.dis_model(fake_images_aug, fake_labels)
                                cls_bcr_real_loss = self.l2_loss(cls_embed_real, cls_embed_real_aug)
                                cls_bcr_fake_loss = self.l2_loss(cls_embed_fake, cls_embed_fake_aug)
                            elif self.conditional_strategy == "ProjGAN_adv":
                                dis_out_real_aug = self.dis_model(real_images_aug, real_labels)
                                dis_out_fake_aug = self.dis_model(fake_images_aug, fake_labels)
                            else:
                                raise NotImplementedError

                            bcr_real_loss = self.l2_loss(dis_out_real, dis_out_real_aug)
                            bcr_fake_loss = self.l2_loss(dis_out_fake, dis_out_fake_aug)
                            if self.conditional_strategy in ["ACGAN", "NT_Xent_GAN", "Proxy_NCA_GAN", "ContraGAN"]:
                                bcr_real_loss += cls_bcr_real_loss
                                bcr_fake_loss += cls_bcr_fake_loss
                            dis_acml_loss += self.real_lambda*bcr_real_loss + self.fake_lambda*bcr_fake_loss

                        if self.zcr:
                            fake_images_zaug = self.gen_model(zs_t, fake_labels)
                            if self.conditional_strategy == "ACGAN":
                                cls_out_fake_zaug, dis_out_fake_zaug = self.dis_model(fake_images_zaug, fake_labels)
                                cls_zcr_dis_loss = self.l2_loss(cls_out_fake, cls_out_fake_zaug)
                            elif self.conditional_strategy == "ProjGAN" or self.conditional_strategy == "no":
                                dis_out_fake_zaug = self.dis_model(fake_images_zaug, fake_labels)
                            elif self.conditional_strategy in ["ContraGAN", "Proxy_NCA_GAN", "NT_Xent_GAN"]:
                                cls_proxies_fake_zaug, cls_embed_fake_zaug, dis_out_fake_zaug = self.dis_model(fake_images_zaug, fake_labels)
                                cls_zcr_dis_loss = self.l2_loss(cls_embed_fake, cls_embed_fake_zaug)
                            elif self.conditional_strategy == "ProjGAN_adv":
                                dis_out_fake_zaug = self.dis_model(fake_images_zaug, fake_labels)
                            else:
                                raise NotImplementedError

                            zcr_dis_loss = self.l2_loss(dis_out_fake, dis_out_fake_zaug)
                            if self.conditional_strategy in ["ACGAN", "NT_Xent_GAN", "Proxy_NCA_GAN", "ContraGAN"]:
                                zcr_dis_loss += cls_zcr_dis_loss
                            dis_acml_loss += self.dis_lambda*zcr_dis_loss

                        if self.gradient_penalty_for_dis:
                            dis_acml_loss += self.gradient_penalty_lambda*calc_derv4gp(self.dis_model, self.conditional_strategy, real_images,
                                                                                fake_images, real_labels, self.default_device)
                        if self.deep_regret_analysis_for_dis:
                            dis_acml_loss += self.regret_penalty_lambda*calc_derv4dra(self.dis_model, self.conditional_strategy, real_images,
                                                                                real_labels, self.default_device)
                        if self.ada:
                            ada_aug_data = torch.tensor((torch.sign(dis_out_real).sum().item(), dis_out_real.shape[0]), device = self.default_device)
                            self.ada_augment += ada_aug_data
                            if self.ada_augment[1] > (self.batch_size*4 - 1):
                                authen_out_signs, num_outputs = self.ada_augment.tolist()
                                r_t_stat = authen_out_signs/num_outputs
                                sign = 1 if r_t_stat > self.ada_target else -1
                                self.ada_aug_p += sign*self.ada_aug_step*num_outputs
                                self.ada_aug_p = min(1.0, max(0.0, self.ada_aug_p))
                                self.ada_augment.mul_(0.0)

                        dis_acml_loss = dis_acml_loss/self.accumulation_steps

                    if self.mixed_precision:
                        self.scaler.scale(dis_acml_loss).backward()
                    else:
                        dis_acml_loss.backward()

                if self.mixed_precision:
                    self.scaler.step(self.D_optimizer)
                    self.scaler.update()
                else:
                    self.D_optimizer.step()

                if self.weight_clipping_for_dis:
                    for p in self.dis_model.parameters():
                        p.data.clamp_(-self.weight_clipping_bound, self.weight_clipping_bound)

            if step_count % self.print_every == 0 and step_count !=0 and self.logger:
                if self.d_spectral_norm:
                    dis_sigmas = calculate_all_sn(self.dis_model)
                    self.writer.add_scalars('SN_of_dis', dis_sigmas, step_count)

            # ================== TRAIN G ================== #
            toggle_grad(self.dis_model, False, freeze_layers=-1)
            toggle_grad(self.gen_model, True, freeze_layers=-1)
            for step_index in range(self.g_steps_per_iter):
                self.G_optimizer.zero_grad()
                for acml_step in range(self.accumulation_steps):
                    with torch.cuda.amp.autocast() if self.mixed_precision else dummy_context_mgr() as mpc:
                        if self.zcr:
                            zs, fake_labels, zs_t = sample_latents(self.prior, self.batch_size, self.z_dim, 1, self.num_classes,
                                                                   self.sigma_noise, self.default_device)
                        else:
                            zs, fake_labels = sample_latents(self.prior, self.batch_size, self.z_dim, 1, self.num_classes,
                                                             None, self.default_device)
                        if self.latent_op:
                            zs, transport_cost = latent_optimise(zs, fake_labels, self.gen_model, self.dis_model, self.conditional_strategy,
                                                                 self.latent_op_step, self.latent_op_rate, self.latent_op_alpha,
                                                                 self.latent_op_beta, True, self.default_device)
                        if not self.conditional_strategy == 'ProjGAN_adv':
                            fake_images = self.gen_model(zs, fake_labels)
                        else:
                            gen_out_prefc, labels_prefc = self.gen_model(zs, fake_labels, only_l1=True)

                            loss_fake = lambda x: -torch.mean(x)
                            gen_out_adv = PGD_G(gen_out_prefc, labels_prefc, fake_labels, loss_fake, self.gen_model, self.dis_model, steps=self.steps, gamma=self.gamma)

                            fake_images = self.gen_model(gen_out_prefc, labels_prefc, l1=False)
                            fake_images_adv = self.gen_model(gen_out_adv, labels_prefc, l1=False)

                        if self.diff_aug:
                            fake_images = DiffAugment(fake_images, policy=self.policy)
                        if self.ada:
                            fake_images, _ = augment(fake_images, self.ada_aug_p)

                        if self.conditional_strategy == "ACGAN":
                            cls_out_fake, dis_out_fake = self.dis_model(fake_images, fake_labels)
                        elif self.conditional_strategy == "ProjGAN" or self.conditional_strategy == "no":
                            dis_out_fake = self.dis_model(fake_images, fake_labels)
                        elif self.conditional_strategy in ["NT_Xent_GAN", "Proxy_NCA_GAN", "ContraGAN"]:
                            fake_cls_mask = make_mask(fake_labels, self.num_classes, self.default_device)
                            cls_proxies_fake, cls_embed_fake, dis_out_fake = self.dis_model(fake_images, fake_labels)
                        elif self.conditional_strategy == 'ProjGAN_adv':
                            dis_out_fake = self.dis_model(fake_images, fake_labels)
                            dis_out_adv = self.dis_model(fake_images_adv, fake_labels)

                        else:
                            raise NotImplementedError

                        gen_acml_loss = self.G_loss(dis_out_fake)
                        if self.latent_op:
                            gen_acml_loss += transport_cost*self.latent_norm_reg_weight

                        if self.zcr:
                            fake_images_zaug = self.gen_model(zs_t, fake_labels)
                            zcr_gen_loss = -1 * self.l2_loss(fake_images, fake_images_zaug)
                            gen_acml_loss += self.gen_lambda*zcr_gen_loss

                        if self.conditional_strategy == "ACGAN":
                            gen_acml_loss += self.ce_loss(cls_out_fake, fake_labels)
                        elif self.conditional_strategy == "ContraGAN":
                            gen_acml_loss += self.contrastive_lambda*self.contrastive_criterion(cls_embed_fake, cls_proxies_fake, fake_cls_mask, fake_labels, t, self.margin)
                        elif self.conditional_strategy == "Proxy_NCA_GAN":
                            gen_acml_loss += self.contrastive_lambda*self.NCA_criterion(cls_embed_fake, cls_proxies_fake, fake_labels)
                        elif self.conditional_strategy == "NT_Xent_GAN":
                            fake_images_aug = CR_DiffAug(fake_images)
                            _, cls_embed_fake_aug, dis_out_fake_aug = self.dis_model(fake_images_aug, fake_labels)
                            gen_acml_loss += self.contrastive_lambda*self.NT_Xent_criterion(cls_embed_fake, cls_embed_fake_aug, t)
                        elif self.conditional_strategy == 'ProjGAN_adv':
                            gen_acml_loss = (self.G_loss(dis_out_fake) + self.G_loss(dis_out_adv)) / 2
                        else:
                            pass

                        gen_acml_loss = gen_acml_loss/self.accumulation_steps

                    if self.mixed_precision:
                        self.scaler.scale(gen_acml_loss).backward()
                    else:
                        gen_acml_loss.backward()

                if self.mixed_precision:
                    self.scaler.step(self.G_optimizer)
                    self.scaler.update()
                else:
                    self.G_optimizer.step()

                # if ema is True: we update parameters of the Gen_copy in adaptive way.
                if self.ema:
                    self.Gen_ema.update(step_count)

                step_count += 1

            if step_count % self.print_every == 0 and self.logger:
                log_message = LOG_FORMAT.format(step=step_count,
                                                progress=step_count/total_step,
                                                elapsed=elapsed_time(self.start_time),
                                                temperature=t,
                                                ada_p=self.ada_aug_p,
                                                dis_loss=dis_acml_loss.item(),
                                                gen_loss=gen_acml_loss.item(),
                                                )
                self.logger.info(log_message)

                if self.g_spectral_norm:
                    gen_sigmas = calculate_all_sn(self.gen_model)
                    self.writer.add_scalars('SN_of_gen', gen_sigmas, step_count)

                self.writer.add_scalars('Losses', {'discriminator': dis_acml_loss.item(),
                                                   'generator': gen_acml_loss.item()}, step_count)
                if self.ada:
                    self.writer.add_scalar('ada_p', self.ada_aug_p, step_count)

            if step_count % self.save_every == 0 or step_count == total_step:
                if self.evaluate:
                    is_best = self.evaluation(step_count, False, "N/A")
                    self.save(step_count, is_best)
                else:
                    self.save(step_count, False)
        return step_count-1
    ################################################################################################################################

    ################################################################################################################################
    def save(self, step, is_best):
        when = "best" if is_best is True else "current"
        self.dis_model.eval()
        self.gen_model.eval()
        if self.Gen_copy is not None:
            self.Gen_copy.eval()

        if isinstance(self.gen_model, DataParallel):
            gen = self.gen_model.module
            dis = self.dis_model.module
            if self.Gen_copy is not None:
                gen_copy = self.Gen_copy.module
        else:
            gen, dis = self.gen_model, self.dis_model
            if self.Gen_copy is not None:
                gen_copy = self.Gen_copy

        g_states = {'seed': self.train_config['seed'], 'run_name': self.run_name, 'step': step, 'best_step': self.best_step,
                    'state_dict': gen.state_dict(), 'optimizer': self.G_optimizer.state_dict(), 'ada_p': self.ada_aug_p}

        d_states = {'seed': self.train_config['seed'], 'run_name': self.run_name, 'step': step, 'best_step': self.best_step,
                    'state_dict': dis.state_dict(), 'optimizer': self.D_optimizer.state_dict(), 'ada_p': self.ada_aug_p,
                    'best_fid': self.best_fid, 'best_fid_checkpoint_path': self.checkpoint_dir}

        if len(glob.glob(join(self.checkpoint_dir,"model=G-{when}-weights-step*.pth".format(when=when)))) >= 1:
            find_and_remove(glob.glob(join(self.checkpoint_dir,"model=G-{when}-weights-step*.pth".format(when=when)))[0])
            find_and_remove(glob.glob(join(self.checkpoint_dir,"model=D-{when}-weights-step*.pth".format(when=when)))[0])

        g_checkpoint_output_path = join(self.checkpoint_dir, "model=G-{when}-weights-step={step}.pth".format(when=when, step=str(step)))
        d_checkpoint_output_path = join(self.checkpoint_dir, "model=D-{when}-weights-step={step}.pth".format(when=when, step=str(step)))

        if when == "best":
            if len(glob.glob(join(self.checkpoint_dir,"model=G-current-weights-step*.pth".format(when=when)))) >= 1:
                find_and_remove(glob.glob(join(self.checkpoint_dir,"model=G-current-weights-step*.pth".format(when=when)))[0])
                find_and_remove(glob.glob(join(self.checkpoint_dir,"model=D-current-weights-step*.pth".format(when=when)))[0])

            g_checkpoint_output_path_ = join(self.checkpoint_dir, "model=G-current-weights-step={step}.pth".format(when=when, step=str(step)))
            d_checkpoint_output_path_ = join(self.checkpoint_dir, "model=D-current-weights-step={step}.pth".format(when=when, step=str(step)))

            torch.save(g_states, g_checkpoint_output_path_)
            torch.save(d_states, d_checkpoint_output_path_)

        torch.save(g_states, g_checkpoint_output_path)
        torch.save(d_states, d_checkpoint_output_path)

        if self.Gen_copy is not None:
            g_ema_states = {'state_dict': gen_copy.state_dict()}
            if len(glob.glob(join(self.checkpoint_dir, "model=G_ema-{when}-weights-step*.pth".format(when=when)))) >= 1:
                find_and_remove(glob.glob(join(self.checkpoint_dir, "model=G_ema-{when}-weights-step*.pth".format(when=when)))[0])

            g_ema_checkpoint_output_path = join(self.checkpoint_dir, "model=G_ema-{when}-weights-step={step}.pth".format(when=when, step=str(step)))

            if when == "best":
                if len(glob.glob(join(self.checkpoint_dir,"model=G_ema-current-weights-step*.pth".format(when=when)))) >= 1:
                    find_and_remove(glob.glob(join(self.checkpoint_dir,"model=G_ema-current-weights-step*.pth".format(when=when)))[0])

                g_ema_checkpoint_output_path_ = join(self.checkpoint_dir, "model=G_ema-current-weights-step={step}.pth".format(when=when, step=str(step)))

                torch.save(g_ema_states, g_ema_checkpoint_output_path_)

            torch.save(g_ema_states, g_ema_checkpoint_output_path)

        if self.logger:
            self.logger.info("Saved model to {}".format(self.checkpoint_dir))

        self.dis_model.train()
        self.gen_model.train()
        if self.Gen_copy is not None:
            self.Gen_copy.train()
    ################################################################################################################################

    ################################################################################################################################
    def evaluation(self, step, standing_statistics, standing_step):
        with torch.no_grad() if self.latent_op is False else dummy_context_mgr() as mpc:
            self.logger.info("Start Evaluation ({step} Step): {run_name}".format(step=step, run_name=self.run_name))
            is_best = False
            num_split, num_run4PR, num_cluster4PR, beta4PR = 1, 10, 20, 8

            self.dis_model.eval()
            generator = change_generator_mode(self.gen_model, self.Gen_copy, standing_statistics, standing_step, self.prior,
                                              self.batch_size, self.z_dim, self.num_classes, self.default_device, training=False)

            fid_score, self.m1, self.s1 = calculate_fid_score(self.eval_dataloader, generator, self.dis_model, self.inception_model, self.num_eval[self.eval_type],
                                                              self.truncated_factor, self.prior, self.latent_op, self.latent_op_step4eval, self.latent_op_alpha,
                                                              self.latent_op_beta, self.default_device, self.mu, self.sigma, self.run_name)

            kl_score, kl_std = calculate_incep_score(self.eval_dataloader, generator, self.dis_model, self.inception_model, self.num_eval[self.eval_type],
                                                     self.truncated_factor, self.prior, self.latent_op, self.latent_op_step4eval, self.latent_op_alpha,
                                                     self.latent_op_beta, num_split, self.default_device)

            precision, recall, f_beta, f_beta_inv = calculate_f_beta_score(self.eval_dataloader, generator, self.dis_model, self.inception_model, self.num_eval[self.eval_type],
                                                                           num_run4PR, num_cluster4PR, beta4PR, self.truncated_factor, self.prior, self.latent_op,
                                                                           self.latent_op_step4eval, self.latent_op_alpha, self.latent_op_beta, self.default_device)
            PR_Curve = plot_pr_curve(precision, recall, self.run_name, self.logger)
            '''
            if self.D_loss.__name__ != "loss_wgan_dis":
                real_train_acc, fake_acc = calculate_accuracy(self.train_dataloader, generator, self.dis_model, self.D_loss, self.num_eval[self.eval_type],
                                                              self.truncated_factor, self.prior, self.latent_op, self.latent_op_step, self.latent_op_alpha,
                                                              self.latent_op_beta, self.default_device, cr=self.cr, eval_generated_sample=True)

                if self.eval_type == 'train':
                    acc_dict = {'real_train': real_train_acc, 'fake': fake_acc}
                else:
                    real_eval_acc = calculate_accuracy(self.eval_dataloader, generator, self.dis_model, self.D_loss, self.num_eval[self.eval_type],
                                                       self.truncated_factor, self.prior, self.latent_op, self.latent_op_step, self.latent_op_alpha,
                                                       self. latent_op_beta, self.default_device, cr=self.cr, eval_generated_sample=False)
                    acc_dict = {'real_train': real_train_acc, 'real_valid': real_eval_acc, 'fake': fake_acc}

                self.writer.add_scalars('{}/Accuracy'.format(self.prune_round), acc_dict, step)
            '''
            if self.best_fid is None:
                self.best_fid, self.best_step, is_best, f_beta_best, f_beta_inv_best = fid_score, step, True, f_beta, f_beta_inv
            else:
                if fid_score <= self.best_fid:
                    self.best_fid, self.best_step, is_best, f_beta_best, f_beta_inv_best = fid_score, step, True, f_beta, f_beta_inv

            self.writer.add_scalars('FID score', {'using {type} moments'.format(type=self.eval_type):fid_score}, step)
            self.writer.add_scalars('F_beta score', {'{num} generated images'.format(num=str(self.num_eval[self.eval_type])):f_beta}, step)
            self.writer.add_scalars('F_beta_inv score', {'{num} generated images'.format(num=str(self.num_eval[self.eval_type])):f_beta_inv}, step)
            self.writer.add_scalars('IS score', {'{num} generated images'.format(num=str(self.num_eval[self.eval_type])):kl_score}, step)
            self.writer.add_figure('PR_Curve', PR_Curve, global_step=step)
            self.logger.info('F_{beta} score (Step: {step}, Using {type} images): {F_beta}'.format(beta=beta4PR, step=step, type=self.eval_type, F_beta=f_beta))
            self.logger.info('F_1/{beta} score (Step: {step}, Using {type} images): {F_beta_inv}'.format(beta=beta4PR, step=step, type=self.eval_type, F_beta_inv=f_beta_inv))
            self.logger.info('FID score (Step: {step}, Using {type} moments): {FID}'.format(step=step, type=self.eval_type, FID=fid_score))
            self.logger.info('Inception score (Step: {step}, {num} generated images): {IS}'.format(step=step, num=str(self.num_eval[self.eval_type]), IS=kl_score))
            if self.train:
                self.logger.info('Best FID score (Step: {step}, Using {type} moments): {FID}'.format(step=self.best_step, type=self.eval_type, FID=self.best_fid))

            self.dis_model.train()
            generator = change_generator_mode(self.gen_model, self.Gen_copy, standing_statistics, standing_step, self.prior,
                                              self.batch_size, self.z_dim, self.num_classes, self.default_device, training=True)

        return is_best
    ################################################################################################################################

    ################################################################################################################################
    def save_images(self, is_generate, standing_statistics, standing_step, png=True, npz=True):
        with torch.no_grad() if self.latent_op is False else dummy_context_mgr() as mpc:
            self.dis_model.eval()
            generator = change_generator_mode(self.gen_model, self.Gen_copy, standing_statistics, standing_step, self.prior,
                                              self.batch_size, self.z_dim, self.num_classes, self.default_device, training=False)

            if png:
                save_images_png(self.run_name, self.eval_dataloader, self.num_eval[self.eval_type], self.num_classes, generator,
                                self.dis_model, is_generate, self.truncated_factor, self.prior, self.latent_op, self.latent_op_step,
                                self.latent_op_alpha, self.latent_op_beta, self.default_device)
            if npz:
                save_images_npz(self.run_name, self.eval_dataloader, self.num_eval[self.eval_type], self.num_classes, generator,
                                self.dis_model, is_generate, self.truncated_factor, self.prior, self.latent_op, self.latent_op_step,
                                self.latent_op_alpha, self.latent_op_beta, self.default_device)
    ################################################################################################################################

    ################################################################################################################################
    def run_image_visualization(self, nrow, ncol, standing_statistics, standing_step):
        self.logger.info('Start visualizing images....')
        with torch.no_grad() if self.latent_op is False else dummy_context_mgr() as mpc:
            generator = change_generator_mode(self.gen_model, self.Gen_copy, standing_statistics, standing_step, self.prior,
                                              self.batch_size, self.z_dim, self.num_classes, self.default_device, training=False)

            sampler = "default" if self.conditional_strategy == "no" else "class_order_some"
            if self.zcr:
                zs, fake_labels, zs_t = sample_latents(self.prior, self.batch_size, self.z_dim, 1, self.num_classes,
                                                     self.sigma_noise, self.default_device, sampler=sampler)
            else:
                zs, fake_labels = sample_latents(self.prior, self.batch_size, self.z_dim, 1, self.num_classes,
                                                None, self.default_device, sampler=sampler)

            if self.latent_op:
                zs = latent_optimise(zs, fake_labels, self.gen_model, self.dis_model, self.conditional_strategy,
                                        self.latent_op_step, self.latent_op_rate, self.latent_op_alpha, self.latent_op_beta,
                                        False, self.default_device)

            generated_images = generator(zs, fake_labels, evaluation=True)

            plot_img_canvas((generated_images.detach().cpu()+1)/2, "./figures/{run_name}/generated_canvas.png".\
                            format(run_name=self.run_name), self.logger, ncol)

            generator = change_generator_mode(self.gen_model, self.Gen_copy, standing_statistics, standing_step, self.prior,
                                              self.batch_size, self.z_dim, self.num_classes, self.default_device, training=True)
    ################################################################################################################################

    ################################################################################################################################
    def run_linear_interpolation(self, nrow, ncol, fix_z, fix_y, standing_statistics, standing_step):
        self.logger.info('Start linear interpolation analysis....')
        with torch.no_grad() if self.latent_op is False else dummy_context_mgr() as mpc:
            generator = change_generator_mode(self.gen_model, self.Gen_copy, standing_statistics, standing_step, self.prior,
                                              self.batch_size, self.z_dim, self.num_classes, self.default_device, training=False)
            shared = generator.module.shared if isinstance(generator, DataParallel) else generator.shared
            assert int(fix_z)*int(fix_y) != 1, "unable to switch fix_z and fix_y on together!"

            if fix_z:
                zs = torch.randn(nrow, 1, self.z_dim, device=self.default_device)
                zs = zs.repeat(1, ncol, 1).view(-1, self.z_dim)
                name = "fix_z"
            else:
                zs = interp(torch.randn(nrow, 1, self.z_dim, device=self.default_device),
                            torch.randn(nrow, 1, self.z_dim, device=self.default_device),
                            ncol - 2).view(-1, self.z_dim)

            if fix_y:
                ys = sample_1hot(nrow, self.num_classes, device=self.default_device)
                ys = shared(ys).view(nrow, 1, -1)
                ys = ys.repeat(1, ncol, 1).view(nrow * (ncol), -1)
                name = "fix_y"
            else:
                ys = interp(shared(sample_1hot(nrow, self.num_classes)).view(nrow, 1, -1),
                            shared(sample_1hot(nrow, self.num_classes)).view(nrow, 1, -1),
                            ncol-2).view(nrow * (ncol), -1)

            interpolated_images = generator(zs, None, shared_label=ys, evaluation=True)

            plot_img_canvas((interpolated_images.detach().cpu()+1)/2, "./figures/{run_name}/Interpolated_images_{fix_flag}.png".\
                            format(run_name=self.run_name, fix_flag=name), self.logger, ncol)

            generator = change_generator_mode(self.gen_model, self.Gen_copy, standing_statistics, standing_step, self.prior,
                                              self.batch_size, self.z_dim, self.num_classes, self.default_device, training=True)
    ################################################################################################################################

    ################################################################################################################################
    def run_nearest_neighbor(self, nrow, ncol, standing_statistics, standing_step):
        self.logger.info('Start nearest neighbor analysis....')
        with torch.no_grad() if self.latent_op is False else dummy_context_mgr() as mpc:
            generator = change_generator_mode(self.gen_model, self.Gen_copy, standing_statistics, standing_step, self.prior,
                                              self.batch_size, self.z_dim, self.num_classes, self.default_device, training=False)

            resnet50_model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet50', pretrained=True)
            resnet50_conv = nn.Sequential(*list(resnet50_model.children())[:-1]).to(self.default_device)
            if self.n_gpus > 1:
                resnet50_conv = DataParallel(resnet50_conv, output_device=self.default_device)
            resnet50_conv.eval()

            for c in tqdm(range(self.num_classes)):
                fake_images, fake_labels = generate_images_for_KNN(self.batch_size, c, generator, self.dis_model, self.truncated_factor, self.prior, self.latent_op,
                                                                   self.latent_op_step, self.latent_op_alpha, self.latent_op_beta, self.default_device)
                fake_image = torch.unsqueeze(fake_images[0], dim=0)
                fake_anchor_embedding = torch.squeeze(resnet50_conv((fake_image+1)/2))

                num_samples, target_sampler = target_class_sampler(self.train_dataset, c)
                train_dataloader = torch.utils.data.DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=False, sampler=target_sampler,
                                                               num_workers=self.train_config['num_workers'], pin_memory=True)
                train_iter = iter(train_dataloader)
                for batch_idx in range(num_samples//self.batch_size):
                    real_images, real_labels = next(train_iter)
                    real_images = real_images.to(self.default_device)
                    real_embeddings = torch.squeeze(resnet50_conv((real_images+1)/2))
                    if batch_idx == 0:
                        distances = torch.square(real_embeddings - fake_anchor_embedding).mean(dim=1).detach().cpu().numpy()
                        holder = real_images.detach().cpu().numpy()
                    else:
                        distances = np.concatenate([distances, torch.square(real_embeddings - fake_anchor_embedding).mean(dim=1).detach().cpu().numpy()], axis=0)
                        holder = np.concatenate([holder, real_images.detach().cpu().numpy()], axis=0)

                nearest_indices = (-distances).argsort()[-(ncol-1):][::-1]
                if c % nrow == 0:
                    canvas = np.concatenate([fake_image.detach().cpu().numpy(), holder[nearest_indices]], axis=0)
                elif c % nrow == nrow-1:
                    row_images = np.concatenate([fake_image.detach().cpu().numpy(), holder[nearest_indices]], axis=0)
                    canvas = np.concatenate((canvas, row_images), axis=0)
                    plot_img_canvas((torch.from_numpy(canvas)+1)/2, "./figures/{run_name}/Fake_anchor_{ncol}NN_{cls}.png".\
                                    format(run_name=self.run_name,ncol=ncol, cls=c), self.logger, ncol)
                else:
                    row_images = np.concatenate([fake_image.detach().cpu().numpy(), holder[nearest_indices]], axis=0)
                    canvas = np.concatenate((canvas, row_images), axis=0)

            generator = change_generator_mode(self.gen_model, self.Gen_copy, standing_statistics, standing_step, self.prior,
                                              self.batch_size, self.z_dim, self.num_classes, self.default_device, training=True)
    ################################################################################################################################

    ################################################################################################################################
    def run_frequency_analysis(self, num_images, standing_statistics, standing_step):
        self.logger.info('Start linear interpolation analysis....')
        with torch.no_grad() if self.latent_op is False else dummy_context_mgr() as mpc:
            generator = change_generator_mode(self.gen_model, self.Gen_copy, standing_statistics, standing_step, self.prior,
                                              self.batch_size, self.z_dim, self.num_classes, self.default_device, training=False)

            train_iter = iter(self.train_dataloader)
            num_batches = num_images//self.batch_size
            for i in range(num_batches):
                if self.zcr:
                    zs, fake_labels, zs_t = sample_latents(self.prior, self.batch_size, self.z_dim, 1, self.num_classes,
                                                           self.sigma_noise, self.default_device)
                else:
                    zs, fake_labels = sample_latents(self.prior, self.batch_size, self.z_dim, 1, self.num_classes,
                                                     None, self.default_device)

                if self.latent_op:
                    zs = latent_optimise(zs, fake_labels, self.gen_model, self.dis_model, self.conditional_strategy,
                                         self.latent_op_step, self.latent_op_rate, self.latent_op_alpha, self.latent_op_beta,
                                         False, self.default_device)

                real_images, real_labels = next(train_iter)
                fake_images = generator(zs, fake_labels, evaluation=True).detach().cpu().numpy()

                real_images = np.asarray((real_images + 1)*127.5, np.uint8)
                fake_images = np.asarray((fake_images + 1)*127.5, np.uint8)

                if i == 0:
                    real_array = real_images
                    fake_array = fake_images
                else:
                    real_array = np.concatenate([real_array, real_images], axis = 0)
                    fake_array = np.concatenate([fake_array, fake_images], axis = 0)

            N, C, H, W = np.shape(real_array)
            real_r, real_g, real_b = real_array[:,0,:,:], real_array[:,1,:,:], real_array[:,2,:,:]
            real_gray = 0.2989 * real_r + 0.5870 * real_g + 0.1140 * real_b
            fake_r, fake_g, fake_b = fake_array[:,0,:,:], fake_array[:,1,:,:], fake_array[:,2,:,:]
            fake_gray = 0.2989 * fake_r + 0.5870 * fake_g + 0.1140 * fake_b
            for j in tqdm(range(N)):
                real_gray_f = np.fft.fft2(real_gray[j] - ndimage.median_filter(real_gray[j], size= H//8))
                fake_gray_f = np.fft.fft2(fake_gray[j] - ndimage.median_filter(fake_gray[j], size=H//8))

                real_gray_f_shifted = np.fft.fftshift(real_gray_f)
                fake_gray_f_shifted = np.fft.fftshift(fake_gray_f)

                if j == 0:
                    real_gray_spectrum = 20*np.log(np.abs(real_gray_f_shifted))/N
                    fake_gray_spectrum = 20*np.log(np.abs(fake_gray_f_shifted))/N
                else:
                    real_gray_spectrum += 20*np.log(np.abs(real_gray_f_shifted))/N
                    fake_gray_spectrum += 20*np.log(np.abs(fake_gray_f_shifted))/N

            plot_spectrum_image(real_gray_spectrum, fake_gray_spectrum, self.run_name, self.logger)

            generator = change_generator_mode(self.gen_model, self.Gen_copy, standing_statistics, standing_step, self.prior,
                                              self.batch_size, self.z_dim, self.num_classes, self.default_device, training=True)
    ################################################################################################################################

def PGD(x, label, loss, model=None, steps=1, gamma=0.1, eps=(1/255), randinit=False, clip=False):

    # Compute loss
    x_adv = x.clone()
    if randinit:
        # adv noise (-eps, eps)
        x_adv += (2.0 * torch.rand(x_adv.shape).cuda() - 1.0) * eps
    x_adv = x_adv.cuda()
    x = x.cuda()

    for t in range(steps):
        out = model(x_adv, label, only_fc=True)
        loss_adv0 = -loss(out)
        grad0 = torch.autograd.grad(loss_adv0, x_adv, only_inputs=True)[0]
        x_adv.data.add_(gamma * torch.sign(grad0.data))

        if clip:
            linfball_proj(x, eps, x_adv, in_place=True)

    return x_adv

def PGD_G(x, gen_labels, label, loss, gen_model, dis_model, steps=1, gamma=0.1, eps=(1/255), randinit=False, clip=False):

    # Compute loss
    x_adv = x.clone()
    x_adv = x_adv.cuda()
    x = x.cuda()

    for t in range(steps):
        out = gen_model(x_adv, gen_labels, l1=False)
        out = dis_model(out, label)
        loss_adv0 = -loss(out)
        grad0 = torch.autograd.grad(loss_adv0, x_adv, only_inputs=True)[0]
        x_adv.data.add_(gamma * torch.sign(grad0.data))

        if clip:
            linfball_proj(x, eps, x_adv, in_place=True)

    return x_adv


import os
import wave

import numpy as np
import pytest

from pydrobert.speech import command_line

def test_compute_feats_from_kaldi_tables(temp_dir):
    kaldi_io = pytest.importorskip("pydrobert.kaldi.io")
    np.random.seed(5)
    wav_scp = os.path.join(temp_dir, "wav.scp")
    feat_ark = os.path.join(temp_dir, "feat.ark")
    wav_file_dir = os.path.join(temp_dir, "wav")
    num_utts = 100
    max_samples = 32000
    # sampling_rate = 16000
    num_digits = int(np.log10(num_utts)) + 1
    utt_fmt = "utt{{:0{}d}}".format(num_digits)
    computer_json_path = os.path.join(
        os.path.dirname(os.path.realpath(__file__)), "data", "fbank.json"
    )
    preprocessor_json_path = os.path.join(temp_dir, "preprocess.json")
    postprocessor_json_path = os.path.join(temp_dir, "unit.json")
    with open(preprocessor_json_path, "w") as f:
        f.write('["dither"]\n')
    with open(postprocessor_json_path, "w") as f:
        f.write('["unit"]\n')
    if not os.path.isdir(wav_file_dir):
        os.makedirs(wav_file_dir)
    with open(wav_scp, "w") as scp:
        for utt_idx in range(num_utts):
            utt_id = utt_fmt.format(utt_idx)
            wav_path = os.path.join(wav_file_dir, utt_id + ".wav")
            num_samples = np.random.randint(max_samples)
            signal = np.random.randint(
                -(2 ** 15), 2 ** 15 - 1, num_samples, dtype=np.int16
            )
            wv = wave.open(wav_path, "wb")
            wv.setnchannels(1)
            wv.setsampwidth(2)
            wv.setframerate(16000)
            wv.writeframes(signal.tobytes())
            wv.close()
            scp.write("{} {}\n".format(utt_id, wav_path))
    args = [
        "scp,s:" + wav_scp,
        "ark:" + feat_ark,
        computer_json_path,
        "--seed=30",
        "--preprocess={}".format(preprocessor_json_path),
        "--postprocess={}".format(postprocessor_json_path),
    ]
    assert not command_line.compute_feats_from_kaldi_tables(args)
    exps = []
    with kaldi_io.open("ark,s:" + feat_ark, "bm") as feats:
        not_all_feats = 0
        for utt_idx, feat in enumerate(feats):
            not_all_feats |= feat.shape[0]
            assert feat.shape[0] == 0 or feat.shape[-1] == 40
            exps.append(feat)
        assert utt_idx == num_utts - 1
        assert not_all_feats
    np.random.seed(300)
    assert not command_line.compute_feats_from_kaldi_tables(args)
    with kaldi_io.open("ark,s:" + feat_ark, "bm") as feats:
        for utt_idx, (exp, act) in enumerate(zip(exps, feats)):
            assert np.allclose(act, exp)

def test_signals_to_torch_feat_dir(temp_dir):
    torch = pytest.importorskip("torch")
    torch.manual_seed(50)
    feat_dir = os.path.join(temp_dir, "feat")
    raw_dir = os.path.join(temp_dir, "raw")
    map_path = os.path.join(temp_dir, "map")
    computer_json_path = os.path.join(
        os.path.dirname(os.path.realpath(__file__)), "data", "fbank.json"
    )
    preprocessor_json_path = os.path.join(temp_dir, "preprocess.json")
    with open(preprocessor_json_path, "w") as f:
        f.write('["dither"]\n')
    assert os.path.isfile(computer_json_path)
    num_utts, max_samples = 100, 1600
    num_digits = torch.log10(torch.tensor(float(num_utts))).long().item() + 1
    utt_fmt = "utt{{:0{}d}}".format(num_digits)
    if not os.path.isdir(raw_dir):
        os.makedirs(raw_dir)
    utt_ids = []
    utt2signal = dict()
    with open(map_path, "w") as mp:
        for utt_idx in range(num_utts):
            utt_id = utt_fmt.format(utt_idx)
            utt_ids.append(utt_id)
            num_samples = torch.randint(1, max_samples, (1,)).long().item()
            signal = torch.randint(
                -(2 ** 15), 2 ** 15 - 1, (num_samples,), dtype=torch.float32
            )
            utt2signal[utt_id] = signal
            file_type = torch.randint(3, (1,)).long().item()
            if file_type == 2:  # wave file
                path = os.path.join(raw_dir, "{}.wav".format(utt_idx))
                wv = wave.open(path, "wb")
                wv.setnchannels(1)
                wv.setsampwidth(2)
                wv.setframerate(16000)
                wv.writeframes(signal.to(torch.int16).numpy().tobytes())
                wv.close()  # py2.7 doesn't have context manager
            elif file_type == 1:  # npy file
                path = os.path.join(raw_dir, "{}.npy".format(utt_idx))
                np.save(path, signal.numpy())
            else:  # torch file
                path = os.path.join(raw_dir, "{}.pt".format(utt_idx))
                torch.save(signal, path)
            mp.write("{} {}\n".format(utt_id, path))
    args = [map_path, computer_json_path, feat_dir]
    assert not command_line.signals_to_torch_feat_dir(args)
    for utt_id in utt_ids:
        feat = torch.load(os.path.join(feat_dir, "{}.pt".format(utt_id)))
        assert feat.shape[-1] == 40
        del feat
    args.pop(1)  # remove computer -- storing raw
    torch.manual_seed(30)
    assert not command_line.signals_to_torch_feat_dir(args)
    for utt_id, exp in list(utt2signal.items()):
        act = torch.load(os.path.join(feat_dir, "{}.pt".format(utt_id)))
        assert torch.allclose(exp, act.flatten())
    torch.manual_seed(40)
    args.append("--seed=1")
    args.append("--preprocess={}".format(preprocessor_json_path))
    assert not command_line.signals_to_torch_feat_dir(args)
    for utt_id in utt_ids:
        signal = utt2signal[utt_id]
        noisy = torch.load(os.path.join(feat_dir, "{}.pt".format(utt_id)))
        assert not torch.allclose(signal, noisy.flatten())
        utt2signal[utt_id] = noisy
    torch.manual_seed(70)
    args.append("--num-workers=4")
    assert not command_line.signals_to_torch_feat_dir(args)
    for utt_id, exp in list(utt2signal.items()):
        act = torch.load(os.path.join(feat_dir, "{}.pt".format(utt_id)))
        assert torch.allclose(exp, act)


#!/usr/bin/env python

import sys, os, glob, string
import numpy as np
from tqdm import tqdm
import time
from pyraf import iraf
import odi_config as odi

try:
    object_str, filters, instrument, images, illcor_flag, skyflat_src, wcs_flag, reproject_flag, scale_flag, scale_ref, stack_flag, align_flag, gaia_flag, cluster_flag, ra_center, dec_center, min_radius = odi.cfgparse('config.yaml')
except IOError:
    print('config.yaml does not exist, quitting...')
    exit()

np.warnings.filterwarnings('ignore')

# for basic processing, filter shouldn't matter enough to consider separately
# or rather, this script already handles that just fine
# so just stick all the image names together into one long list
images_ = [img for sublist in list(images.values()) for img in sublist]
# print images_

rad, decd = odi.get_targ_ra_dec(images_[0], 'OTA33.SCI')
if gaia_flag:
    source = 'gaia'
else:
    source = 'sdss'
inst = odi.instrument(instrument)

# start by creating "raw" individual OTA images
for img in images_:
    for key in tqdm(odi.OTA_dictionary, desc='Creating raw OTA images:', ncols=0):
        ota = odi.OTA_dictionary[key]
        ota_raw_data = img.f+'['+ota+']'
        raw_img = 'raw_'+ota+'.'+img.stem()
        if not os.path.isfile(odi.rawpath+raw_img):
            iraf.unlearn(iraf.imutil.imcopy)
            iraf.imutil.imcopy.setParam('input',ota_raw_data)
            iraf.imutil.imcopy.setParam('output',odi.rawpath+raw_img)
            iraf.imutil.imcopy.setParam('verbose','no')
            iraf.imutil.imcopy(mode='h')

# determine guide OTAs
if not os.path.isfile('guide_otas.txt'):
    guide_file = open('guide_otas.txt','w+')
    guide_otas = list() # make an empty list so we can keep track of the guide OTAs
    for img in images_:
        for key in tqdm(odi.OTA_dictionary, desc='Determining Guide OTAs for {:s}'.format(img.stem()), ncols=0):
            ota = odi.OTA_dictionary[key]
            guide = odi.is_guide_ota(img, ota) # this function uses corner/center ratio method
            fullid = ota+'.'+img.stem()
            # tqdm.write("{:} {:}".format(fullid, guide))
            if guide:   # if it looks like a guide OTA, add it to the list
                guide_otas.append(fullid)
                print(fullid, file=guide_file)
    guide_otas = np.array(guide_otas)
else:
    guide_otas = np.loadtxt('guide_otas.txt',usecols=(0,),unpack=True,dtype=str)
    # guide_otas = list(guide_otas_r)

#Create offline catalogs
for img in images_:
    for key in tqdm(odi.OTA_dictionary, desc='Retrieving QR SDSS and Gaia catalogs for {:s}'.format(img.stem()), ncols=0):
        ota = odi.OTA_dictionary[key]
        outputsd = odi.sdsspath+'offline_'+ota+'.'+img.base()+'.sdss'
        if not os.path.isfile(outputsd):
            x,y = odi.get_sdss_coords_offline(img,ota,inst,output=outputsd)
            # output2m = odi.twomasspath+'offline_'+ota+'.'+img.base()+'.mass'
            # x,y = odi.get_2mass_coords_offline(img,ota,inst,output=output2m)
        if gaia_flag == True:
            outputg = odi.gaiapath+'offline_'+ota+'.'+img.base()+'.gaia'
            if not os.path.isfile(outputg):
                if cluster_flag == True:
                    odi.get_gaia_coords(img,ota,inst,
                                        output=outputg,
                                        cluster=cluster_flag,
                                        racenter=float(ra_center),
                                        deccenter=float(dec_center),
                                        min_radius=float(min_radius),
                                        G_lim = 20.75)
                else:
                    odi.get_gaia_coords(img,ota,inst,
                                        output=outputg,
                                        cluster=cluster_flag)

# if illcor_flag:
listfiles = glob.glob('*.lis')
if len(listfiles) == 0:
    odi.imcombine_lists(images_, filters, guide_otas)
else:
    print('imcombine lists done')

if not os.path.isfile('bpms.done'):
    for img in images_:
        print('updating bpms for', img.stem())
        for key in tqdm(odi.OTA_dictionary):
            ota = odi.OTA_dictionary[key]
            odi.make_bpms(img, ota)
    with open('bpms.done', 'w+') as bpm:
        print('bpms are done!', file=bpm)

listfiles = glob.glob(odi.skyflatpath+'*.med.fits')
if len(listfiles) == 0:
    for filter in filters:
        if illcor_flag:
            if len(images_) > len(filters)+1:
                odi.dark_sky_flat(filter, box_size=51)
            else :
                odi.dark_sky_flat(filter, box_size=151)
        else:
            print('not making dark sky flats for', filter)
else:
    print('dark sky flats done')

if not os.path.isfile('derived_props.txt'):
    f1 = open('derived_props.txt','w+')
    print('# img                        ota       filter  guide  fwhm  zp_med  zp_std  bg_mean   bg_med   bg_std', file=f1)
    finished = list()
else:
    fwhm, zp_med, zp_std, bg_mean, bg_median, bg_std = np.loadtxt('derived_props.txt',usecols=(4,5,6,7,8,9),unpack=True)
    imgnum, ota_d, filt_d, guide_d = np.loadtxt('derived_props.txt',usecols=(0,1,2,3),unpack=True,dtype=str)
    finished = list(zip(imgnum,ota_d,filt_d))
    f1 = open('derived_props.txt','a+')

for img in images_:
    otalist = sorted(odi.OTA_dictionary.keys())
    for key in tqdm(otalist):
        ota = odi.OTA_dictionary[key]
        fullid = ota+'.'+img.stem()
        hdulist = odi.fits.open(img.f)
        hdr = hdulist[0].header
        filt = hdr['filter']
        finishcheck = (img.stem(),ota,filt)
        if finishcheck in finished:
            already = 0
            # guide_d = fullid in guide_otas
            # if guide_d:
            #     tqdm.write("{:} {:}".format(fullid, guide_d))
        else:
            # image_to_correct = img.f+'['+ota+']'
            image_to_correct = odi.rawpath+'raw_'+ota+'.'+img.stem()
            correction_image = ota+'.'+filt+'.med.smooth.fits'
            corrected_image = 'illcor_'+ota+'.'+img.stem()
            tqdm.write(corrected_image)
            if not os.path.isfile(odi.illcorpath+corrected_image):
                odi.illumination_corrections(image_to_correct, correction_image, corrected_image, do_correction=illcor_flag)
            gaps = odi.get_gaps(img, ota)
            reprojed_image = 'reproj_'+ota+'.'+img.stem()
            # wcsrefimg = odi.illcorpath+'illcor_OTA33.SCI.'+images_[0].stem()
            # wcsref = odi.fits.getheader(wcsrefimg)
            wcsref = odi.illcorpath+'illcor_OTA33.SCI.'+images_[0].stem()
            if not os.path.isfile(odi.reprojpath+reprojed_image):
                if wcs_flag:
                    pixcrd3 = odi.list_wcs_coords(img, ota, gaps, inst,output=img.nofits()+'.'+ota+'.radec.coo', gmaglim=23., stars_only=True, offline = True, source = source)
                    try:
                        odi.fix_wcs(img, ota, coords=img.nofits()+'.'+ota+'.radec.coo', iters=1)
                    except:
                        try:
                            print('msccmatch failed, wait a second and try again')
                            time.sleep(1.0)
                            odi.fix_wcs(img, ota, coords=img.nofits()+'.'+ota+'.radec.coo', iters=1)
                        except:
                            print('there might be too few stars for msccmatch, just skip it.')
                if reproject_flag:
                    odi.reproject_ota(img, ota, rad, decd, wcsref)
                    odi.tpv2tan_hdr(img, ota)
            gaps = odi.get_gaps_rep(img, ota)
            odi.refetch_sdss_coords(img, ota, gaps, inst,gmaglim=21.5,offline = True,source=source)
            #run an additional refetch to get the xy for 2mass so they can be used for scaling
            # odi.repoxy_offline(img, ota, gaps, inst,gmaglim=21.5,source='twomass')
            odi.repoxy_offline(img, ota, gaps, inst,gmaglim=21.5,source='gaia')
            fwhm = odi.getfwhm_ota(img, ota, gaia=gaia_flag)
            if 'odi_NB695' in filters:
                zp_med, zp_std = 99.99,99.99
            elif source == 'sdss':
                zp_med, zp_std = 99.99,99.99
            elif source == 'twomass':
                zp_med, zp_std = 99.99,99.99
            elif source == 'gaia':
                zp_med, zp_std = 99.99,99.99
            if not os.path.isfile(odi.bgsubpath+'bgsub_'+ota+'.'+img.stem()):
                bg_mean, bg_median, bg_std = odi.bgsub_ota(img, ota, apply=True)
            else:
                bg_mean, bg_median, bg_std = odi.bgsub_ota(img, ota, apply=False)
            guide_d = fullid in guide_otas
            tqdm.write("{:} {:}".format(fullid, guide_d))
            print("{0:s}  {1:9s}  {2:5s}  {3!s:5s}  {4:3.1f}   {5:5.2f}   {6:5.2f}  {7:8.2f}  {8:8.2f} {9:8.2f}".format(img.stem(), ota, filt, guide, fwhm, zp_med, zp_std, bg_mean, bg_median, bg_std), file=f1)
            dim_stats = odi.check_mask_dim(img,ota)
            if not dim_stats:
                print('mask dimensions do not match image')
                print('redo', img, ota)
                raise ValueError
f1.close()


import numpy as np
from keras.models import load_model
from tqdm import tqdm
from sklearn.metrics import classification_report
import keras.backend as K
from keras.preprocessing import sequence

import utils as U
import reader as dataset
from my_layers import Attention, Average, WeightedSum, WeightedAspectEmb, MaxMargin

######### Get hyper-params in order to rebuild the model architecture ###########
# The hyper parameters should be exactly the same as those used for training

parser = U.add_common_args()
args = parser.parse_args()

out_dir = args.out_dir_path + '/' + args.domain
# out_dir = '../pre_trained_model/' + args.domain
U.print_args(args)

# assert args.domain in {'restaurant', 'beer'}

###### Get test data #############
vocab, train_x, _, overall_maxlen = dataset.get_data(args.domain, vocab_size=args.vocab_size, maxlen=args.maxlen)
test_x = train_x
test_x = sequence.pad_sequences(test_x, maxlen=overall_maxlen)
test_length = test_x.shape[0]
splits = []
for i in range(1, test_length // args.batch_size):
    splits.append(args.batch_size * i)
if test_length % args.batch_size:
    splits += [(test_length // args.batch_size) * args.batch_size]
test_x = np.split(test_x, splits)

############# Build model architecture, same as the model used for training #########

## Load the save model parameters
model = load_model(out_dir + '/model_param',
                   custom_objects={"Attention": Attention, "Average": Average, "WeightedSum": WeightedSum,
                                   "MaxMargin": MaxMargin, "WeightedAspectEmb": WeightedAspectEmb,
                                   "max_margin_loss": U.max_margin_loss},
                   compile=True)

################ Evaluation ####################################

def evaluation(true, predict, domain):
    true_label = []
    predict_label = []

    if domain == 'restaurant':

        for line in predict:
            predict_label.append(line.strip())

        for line in true:
            true_label.append(line.strip())

        print(classification_report(true_label, predict_label,
                                    ['Food', 'Staff', 'Ambience', 'Anecdotes', 'Price', 'Miscellaneous'], digits=3))

    elif domain == 'drugs_cadec':
        for line in predict:
            predict_label.append(line.strip())

        for line in true:
            true_label.append(line.strip())

        print(classification_report(true_label, predict_label, digits=3))

    else:
        for line in predict:
            label = line.strip()
            if label == 'smell' or label == 'taste':
                label = 'taste+smell'
            predict_label.append(label)

        for line in true:
            label = line.strip()
            if label == 'smell' or label == 'taste':
                label = 'taste+smell'
            true_label.append(label)

        print(classification_report(true_label, predict_label,
                                    ['feel', 'taste+smell', 'look', 'overall', 'None'], digits=3))

def prediction(test_labels, aspect_probs, cluster_map, domain):
    label_ids = np.argsort(aspect_probs, axis=1)[:, -1]
    predict_labels = [cluster_map[label_id] for label_id in label_ids]
    evaluation(open(test_labels), predict_labels, domain)

## Create a dictionary that map word index to word
vocab_inv = {}
for w, ind in vocab.items():
    vocab_inv[ind] = w

test_fn = K.function([model.get_layer('sentence_input').input, K.learning_phase()],
                     [model.get_layer('att_weights').output, model.get_layer('p_t').output])
att_weights, aspect_probs = [], []
for batch in tqdm(test_x):
    cur_att_weights, cur_aspect_probs = test_fn([batch, 0])
    att_weights.append(cur_att_weights)
    aspect_probs.append(cur_aspect_probs)

att_weights = np.concatenate(att_weights)
aspect_probs = np.concatenate(aspect_probs)

######### Topic weight ###################################

topic_weight_out = open(out_dir + '/topic_weights', 'wt', encoding='utf-8')
labels_out = open(out_dir + '/labels.txt', 'wt', encoding='utf-8')
print('Saving topic weights on test sentences...')
for probs in aspect_probs:
    labels_out.write(str(np.argmax(probs)) + "\n")
    weights_for_sentence = ""
    for p in probs:
        weights_for_sentence += str(p) + "\t"
    weights_for_sentence.strip()
    topic_weight_out.write(weights_for_sentence + "\n")
print(aspect_probs)

## Save attention weights on test sentences into a file
att_out = open(out_dir + '/att_weights', 'wt', encoding='utf-8')
print('Saving attention weights on test sentences...')
test_x = np.concatenate(test_x)
for c in range(len(test_x)):
    att_out.write('----------------------------------------\n')
    att_out.write(str(c) + '\n')

    word_inds = [i for i in test_x[c] if i != 0]
    line_len = len(word_inds)
    weights = att_weights[c]
    weights = weights[(overall_maxlen - line_len):]

    words = [vocab_inv[i] for i in word_inds]
    att_out.write(' '.join(words) + '\n')
    for j in range(len(words)):
        att_out.write(words[j] + ' ' + str(round(weights[j], 3)) + '\n')

######################################################
# Uncomment the below part for F scores
######################################################

## cluster_map need to be specified manually according to the top words in each inferred aspect (save in aspect.log)

# map for the pre-trained restaurant model (under pre_trained_model/restaurant)
# cluster_map = {0: 'Food', 1: 'Miscellaneous', 2: 'Miscellaneous', 3: 'Food',
#            4: 'Miscellaneous', 5: 'Food', 6:'Price',  7: 'Miscellaneous', 8: 'Staff',
#            9: 'Food', 10: 'Food', 11: 'Anecdotes',
#            12: 'Ambience', 13: 'Staff'}

# print '--- Results on %s domain ---' % (args.domain)
# test_labels = '../preprocessed_data/%s/test_label.txt' % (args.domain)
# prediction(test_labels, aspect_probs, cluster_map, domain=args.domain)


import cv2
import face_recognition as FR
from DistanceMetrics import Similarity
import numpy as np
import tensorflow as tf
from facenet import facenet

def prewhiten(image):
    mean = np.mean(image)
    std = np.std(image)
    std_adj = np.maximum(std, 1.0/np.sqrt(image.size))
    processed = np.multiply(np.subtract(image, mean), 1/std_adj)

    return processed

def hist_equalize(image):
    lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(16, 16))
    l1 = clahe.apply(l)
    processed = cv2.merge((l1, a, b))
    processed = cv2.cvtColor(processed, cv2.COLOR_LAB2BGR)

    return processed

cap = cv2.VideoCapture(0)
dist = Similarity()

with tf.Graph().as_default():
    with tf.Session() as session:

        facenet.load_model('20180402-114759.pb')
        img_holder = tf.get_default_graph().get_tensor_by_name(
                                    'input:0')
        embeddings = tf.get_default_graph().get_tensor_by_name(
                                    'embeddings:0')
        phase_train = tf.get_default_graph().get_tensor_by_name(
                                    'phase_train:0')

        test_image = cv2.imread('test2.jpg')
        #test_image = hist_equalize(test_image)
        (y1, x2, y2, x1) = FR.face_locations(test_image, model='hog')[0]
        test_face = cv2.resize(test_image[y1:y2, x1:x2], (160, 160))
        #test_face = hist_equalize(test_face)
        feed_dict = {img_holder:[test_face], phase_train:False}
        test_enc1 = session.run(embeddings, feed_dict=feed_dict)
        cv2.imshow('Test 1', test_face)
        test_face = prewhiten(test_face)
        feed_dict = {img_holder:[test_face], phase_train:False}
        test_enc2 = session.run(embeddings, feed_dict=feed_dict)
        cv2.imshow('Test 2', test_face)
        while True:
            ret, frame = cap.read()
            if not ret:
                continue

            image = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            try:
                (y1, x2, y2, x1) = FR.face_locations(image,
                                      model='hog')[0]
                frame = hist_equalize(frame)
                face = cv2.resize(frame[y1:y2, x1:x2], (160, 160))
                #face = hist_equalize(face)
                cv2.imshow('Unprocessed', face)
                processed = prewhiten(face)
                cv2.imshow('Processed', processed)
            except:
                pass
            try:
                feed_dict = {img_holder:[face], phase_train:False}
                e1 = session.run(embeddings, feed_dict=feed_dict)
                feed_dict = {img_holder:[processed], phase_train:False}
                e2 = session.run(embeddings, feed_dict=feed_dict)
                distance = dist.fractional_distance(e1[0], test_enc1, fraction=0.5)
                print('Distance, not pre-whitened : ', distance)
                distance = dist.fractional_distance(e2[0], test_enc2, fraction=0.5)
                print('Distance, pre-whitened both: ', distance)
                print('')
            except:
                pass

            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                break

cv2.destroyAllWindows()


from __future__ import absolute_import, division, print_function, unicode_literals

import argparse
import logging
import os
import random, sys
from io import open

import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset, RandomSampler
from torch.utils.data.distributed import DistributedSampler

sys.path.append("/u/flashscratch/d/datduong/pytorch-pretrained-BERT")

import pytorch_pretrained_bert.modeling as modeling
import pytorch_pretrained_bert.tokenization
import pytorch_pretrained_bert.optimization

from torch.utils.data import Dataset
import random

PATH="/u/home/w/wenlongx/BERT/pytorch-pretrained-BERT"

# This represents a file that lists all of the possible words occurring in the GO dataset
VOCAB_PATH="/u/flashscratch/d/datduong/goAndGeneAnnotationDec2018/GO_db_complete_vocab.txt"

# This we should not change, because it represents the initial word embeddings that
# Dat has already pretrained on PubMed data
INIT_EMB_PATH="/u/flashscratch/d/datduong/w2vModel1Gram11Nov2017/w2vModel1Gram11Nov2017.txt"

# Pickle of the GO term objects
PICKLE_PATH="/u/flashscratch/d/datduong/goAndGeneAnnotationDec2018/w2vDim300/EntDataJan19w300Base/GO_all_info_go_graph.pickle"

input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])

config = modeling.BertConfig(vocab_size_or_config_json_file=23408,
        hidden_size=300,
        num_hidden_layers=12,
        num_attention_heads=12,
        intermediate_size=3072,
        max_position_embeddings=23408,
        vocab_path=VOCAB_PATH,
        init_emb_path=INIT_EMB_PATH)

model = modeling.BertModel(config=config)

all_encoder_layers, pooled_output = model(input_ids, token_type_ids, input_mask)


'''TD3 is short for Twin Delayed Deep Deterministic Policy Gradients
Designed to mitigate the function approximation error in actor critic, including
    over-estimation bias, variance and bias due to the accumulation of errors because
    of using a deep neural network to approximate the agent's policy and value
    function
Modifications compared to DDPG:
1. delay the policy update by every other step
2. smooth the calculation of target action (by adding a random noise) when learning the critic
3. perform a double Q learning type modication of learning rule for a critic update

Paper of this method:
http://proceedings.mlr.press/v80/fujimoto18a.html
'''

import numpy as np
import torch as T
import torch.nn.functional as F
import torch.nn as nn
import torch.optim as optim
import os

from util import ReplayBuffer, update_single_target_network_parameters, initWB

class CriticNetwork(nn.Module):
    def __init__(self, crt_lr, input_dims, n_actions, fc_dims,
                name='Critic', chkpt_dir='tmp/td3', layerNorm=True):
        super(CriticNetwork, self).__init__()
        self.input_dims = input_dims
        self.n_actions = n_actions
        self.fc_dims = fc_dims
        self.model_dir = os.path.join(chkpt_dir, name)
        self.layerNorm = layerNorm

        # Fully connected layers
        self.fcs = nn.ModuleList()
        self.lns = nn.ModuleList()
        self.n_fcs = len(self.fc_dims)
        for fc_i in range(self.n_fcs):
            # Fully connected layer
            if fc_i == 0:
                fc = nn.Linear(self.input_dims+self.n_actions, self.fc_dims[0])
            else:
                fc = nn.Linear(self.fc_dims[fc_i-1], self.fc_dims[fc_i])
            initWB(fc)
            self.fcs.append(fc)
            # Layer Normalization over a mini-batch of inputs
            ln = nn.LayerNorm(self.fc_dims[fc_i])
            self.lns.append(ln)

        # Output layer
        self.q = nn.Linear(self.fc_dims[-1], 1)
        initWB(self.q)

        self.optimizer = optim.Adam(self.parameters(), lr=crt_lr)
        self.device = T.device('cuda:0' if T.cuda.is_available() else 'cpu')
        self.to(self.device)

    def forward(self, state, action):
        state_action_value = T.cat([state, action], dim=1)
        for fc, ln in zip(self.fcs, self.lns):
            if self.layerNorm:
                state_action_value = F.relu(ln(fc(state_action_value)))
            else:
                state_action_value = F.relu(fc(state_action_value))

        q = self.q(state_action_value)

        return q

    def save_checkpoint(self):
        T.save(self.state_dict(), self.model_dir)

    def load_checkpoint(self):
        self.load_state_dict(T.load(self.model_dir))

class ActorNetwork(nn.Module):
    def __init__(self, act_lr, input_dims, n_actions, fc_dims,
                 name='Actor', chkpt_dir='tmp/td3', layerNorm=True):
        super(ActorNetwork, self).__init__()
        self.input_dims = input_dims
        self.n_actions = n_actions
        self.fc_dims = fc_dims
        self.model_dir = os.path.join(chkpt_dir, name)
        self.layerNorm = layerNorm

        # Fully connected layers
        self.fcs = nn.ModuleList()
        self.lns = nn.ModuleList()
        self.n_fcs = len(self.fc_dims)
        for fc_i in range(self.n_fcs):
            # Fully connected layer
            if fc_i == 0:
                fc = nn.Linear(self.input_dims, self.fc_dims[0])
            else:
                fc = nn.Linear(self.fc_dims[fc_i-1], self.fc_dims[fc_i])
            initWB(fc)
            self.fcs.append(fc)
            # Layer Normalization over a mini-batch of inputs
            ln = nn.LayerNorm(self.fc_dims[fc_i])
            self.lns.append(ln)

        self.mu = nn.Linear(self.fc_dims[-1], self.n_actions)
        initWB(self.mu)

        self.optimizer = optim.Adam(self.parameters(), lr=act_lr)
        self.device = T.device('cuda:0' if T.cuda.is_available() else 'cpu')
        self.to(self.device)

    def forward(self, state):
        for fc, ln in zip(self.fcs, self.lns):
            if self.layerNorm:
                state = F.relu(ln(fc(state)))
            else:
                state = F.relu(fc(state))

        x = T.tanh(self.mu(state))

        return x

    def save_checkpoint(self):
        T.save(self.state_dict(), self.model_dir)

    def load_checkpoint(self):
        self.load_state_dict(T.load(self.model_dir))

class Agent():
    def __init__(self, input_dims, n_actions,layer_sizes,
                 act_lr=0.00001, crt_lr=0.0001, tau=0.001, gamma=0.99,
                 max_size=1000000, batch_size=64, update_actor_interval=2,
                 noise=0.1, noise_targetAct=0.2,
                 chkpt_dir='tmp/td3', name='td3', layerNorm=True):
        self.input_dims = input_dims
        self.n_actions = n_actions
        self.gamma = gamma
        self.tau = tau
        self.max_action = 1
        self.min_action = -1
        self.memory = ReplayBuffer(max_size, self.input_dims, self.n_actions)
        self.batch_size = batch_size
        self.learn_step_cntr = 0
        self.update_actor_iter = update_actor_interval

        self.actor = ActorNetwork(act_lr, self.input_dims, self.n_actions, layer_sizes,
                                name='Actor_'+name, chkpt_dir=chkpt_dir, layerNorm=layerNorm)

        self.critic_1 = CriticNetwork(crt_lr, self.input_dims, self.n_actions, layer_sizes,
                                name='Critic1_'+name, chkpt_dir=chkpt_dir, layerNorm=layerNorm)
        self.critic_2 = CriticNetwork(crt_lr, self.input_dims, self.n_actions, layer_sizes,
                                name='Critic2_'+name, chkpt_dir=chkpt_dir, layerNorm=layerNorm)

        self.target_actor = ActorNetwork(act_lr, self.input_dims, self.n_actions, layer_sizes,
                                name='TargetActor_'+name, chkpt_dir=chkpt_dir, layerNorm=layerNorm)
        self.target_critic_1 = CriticNetwork(crt_lr, self.input_dims, self.n_actions, layer_sizes,
                                name='TargetCritic1_'+name, chkpt_dir=chkpt_dir, layerNorm=layerNorm)
        self.target_critic_2 = CriticNetwork(crt_lr, self.input_dims, self.n_actions, layer_sizes,
                                name='TargetCritic2_'+name, chkpt_dir=chkpt_dir, layerNorm=layerNorm)

        self.noise = noise
        self.noise_targetAct = noise_targetAct
        self.update_network_parameters(tau=1)

    def choose_action(self, observation):
        state = T.tensor(observation, dtype=T.float).to(self.actor.device)
        mu = self.actor.forward(state).to(self.actor.device)
        mu_prime = mu + T.tensor(np.random.normal(scale=self.noise),
                                 dtype=T.float).to(self.actor.device)

        mu_prime = T.clamp(mu_prime, self.min_action, self.max_action)

        return mu_prime.cpu().detach().numpy()

    def remember(self, state, action, reward, new_state, done):
        self.memory.store_transition(state, action, reward, new_state, done)

    def learn(self):
        if self.memory.mem_cntr < self.batch_size:
            return
        state, action, reward, new_state, done = \
            self.memory.sample_buffer(self.batch_size)

        reward = T.tensor(reward, dtype=T.float).to(self.critic_1.device)
        # done = T.tensor(done).to(self.critic_1.device)
        state_ = T.tensor(new_state, dtype=T.float).to(self.critic_1.device)
        state = T.tensor(state, dtype=T.float).to(self.critic_1.device)
        action = T.tensor(action, dtype=T.float).to(self.critic_1.device)

        target_actions = self.target_actor.forward(state_)
        target_actions = target_actions + \
            T.clamp(T.tensor(np.random.normal(
                scale=self.noise_targetAct)), -0.5, 0.5)
        target_actions = T.clamp(target_actions, self.min_action,
                                 self.max_action)

        q1_ = self.target_critic_1.forward(state_, target_actions).view(-1)
        q2_ = self.target_critic_2.forward(state_, target_actions).view(-1)
        # q1_[done] = 0.0   # In building context, the terminal state does not have 0 value
        # q2_[done] = 0.0
        critic_value_ = T.min(q1_, q2_)
        target = reward + self.gamma*critic_value_

        self.critic_1.optimizer.zero_grad()
        self.critic_2.optimizer.zero_grad()
        q1 = self.critic_1.forward(state, action).view(-1)
        q2 = self.critic_2.forward(state, action).view(-1)
        q1_loss = F.mse_loss(target, q1)
        q2_loss = F.mse_loss(target, q2)
        critic_loss = q1_loss + q2_loss
        critic_loss.backward()
        self.critic_1.optimizer.step()
        self.critic_2.optimizer.step()

        self.learn_step_cntr += 1
        # if self.learn_step_cntr % self.update_actor_iter != 0:
        #     return

        self.actor.optimizer.zero_grad()
        actor_q1_loss = self.critic_1.forward(
            state, self.actor.forward(state))  # can also use the mean
        # of actor_q1_loss and actor_q2_loss, but it would be slower and does not really matter
        actor_loss = -T.mean(actor_q1_loss)
        actor_loss.backward()
        self.actor.optimizer.step()

        self.update_network_parameters()

        return critic_loss.item(), actor_loss.item()

    def update_network_parameters(self, tau=None):
        if tau is None:
            tau = self.tau

        updated_actor = update_single_target_network_parameters(
            self.actor, self.target_actor, tau
        )
        updated_critic_1 = update_single_target_network_parameters(
            self.critic_1, self.target_critic_1, tau
        )
        updated_critic_2 = update_single_target_network_parameters(
            self.critic_2, self.target_critic_2, tau
        )

        self.target_actor.load_state_dict(updated_actor)
        self.target_critic_1.load_state_dict(updated_critic_1)
        self.target_critic_2.load_state_dict(updated_critic_2)

    def save_models(self):
        print('.... saving models ....')
        self.actor.save_checkpoint()
        # self.target_actor.save_checkpoint()
        self.critic_1.save_checkpoint()
        self.critic_2.save_checkpoint()
        # self.target_critic_1.save_checkpoint()
        # self.target_critic_2.save_checkpoint()

    def load_models(self):
        print('.... loading models ....')
        self.actor.load_checkpoint()
        # self.target_actor.load_checkpoint()
        self.critic_1.load_checkpoint()
        self.critic_2.load_checkpoint()
        # self.target_critic_1.load_checkpoint()
        # self.target_critic_2.load_checkpoint()


import pickle
import tensorflow as tf
import cv2
import numpy as np
import time
import logging
import argparse

from tensorflow.python.client import timeline

from network_cmu import CmuNetwork
from common import estimate_pose, CocoPairsRender

logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')

config = tf.ConfigProto()
config.gpu_options.allocator_type = 'BFC'
config.gpu_options.per_process_gpu_memory_fraction = 0.95
config.gpu_options.allow_growth = True

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Tensorflow Openpose Inference')
    parser.add_argument('--imgpath', type=str, default='./images/person3.jpg')
    parser.add_argument('--input-width', type=int, default=320)
    parser.add_argument('--input-height', type=int, default=240)
    parser.add_argument('--stage-level', type=int, default=6)
    parser.add_argument('--model', type=str, default='kakao', help='cmu(original) / kakao(faster version)')
    args = parser.parse_args()

    input_node = tf.placeholder(tf.float32, shape=(None, args.input_height, args.input_width, 3), name='image')

    with tf.Session(config=config) as sess:
        if args.model == 'kakao':
            net = KakaoNetwork({'image': input_node}, trainable=False)
            net.load('./models/numpy/fastopenpose_coco_v170729.npy', sess)
        elif args.model == 'cmu':
            net = CmuNetwork({'image': input_node}, trainable=False)
            net.load('./models/numpy/openpose_coco.npy', sess)
        else:
            raise Exception('Invalid Mode.')

        logging.debug('read image+')
        image = cv2.imread(args.imgpath)
        image = cv2.resize(image, (args.input_width, args.input_height))
        image = image.astype(float)
        image -= 128.0
        image /= 128.0

        vec = sess.run(net.get_output(name='concat_stage7'), feed_dict={'image:0': [image]})

        logging.debug('inference+')
        a = time.time()
        pafMat, heatMat = sess.run(
            [
                net.get_output(name='Mconv7_stage{}_L1'.format(args.stage_level)),
                net.get_output(name='Mconv7_stage{}_L2'.format(args.stage_level))
            ], feed_dict={'image:0': [image]}
        )
        logging.info('inference- elapsed_time={}'.format(time.time() - a))
        a = time.time()
        pafMat, heatMat = sess.run(
            [
                net.get_output(name='Mconv7_stage{}_L1'.format(args.stage_level)),
                net.get_output(name='Mconv7_stage{}_L2'.format(args.stage_level))
            ], feed_dict={'image:0': [image]}
        )
        logging.info('inference- elapsed_time={}'.format(time.time() - a))
        a = time.time()

        run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
        run_metadata = tf.RunMetadata()
        pafMat, heatMat = sess.run(
            [
                net.get_output(name='Mconv7_stage{}_L1'.format(args.stage_level)),
                net.get_output(name='Mconv7_stage{}_L2'.format(args.stage_level))
            ], feed_dict={'image:0': [image]}, options=run_options, run_metadata=run_metadata
        )
        logging.info('inference- elapsed_time={}'.format(time.time() - a))

        tl = timeline.Timeline(run_metadata.step_stats)
        ctf = tl.generate_chrome_trace_format()
        with open('timeline.json', 'w') as f:
            f.write(ctf)

        heatMat, pafMat = heatMat[0], pafMat[0]
        heatMat = np.rollaxis(heatMat, 2, 0)
        pafMat = np.rollaxis(pafMat, 2, 0)

        logging.info('pickle data')
        with open('heatmat.pickle', 'wb') as pickle_file:
            pickle.dump(heatMat, pickle_file, pickle.HIGHEST_PROTOCOL)
        with open('pafmat.pickle', 'wb') as pickle_file:
            pickle.dump(pafMat, pickle_file, pickle.HIGHEST_PROTOCOL)

        logging.info('pose+')
        a = time.time()
        humans = estimate_pose(heatMat, pafMat)
        logging.info('pose- elapsed_time={}'.format(time.time() - a))

        # display
        image = cv2.imread(args.imgpath)
        image_h, image_w = image.shape[:2]
        heat_h, heat_w = heatMat[0].shape[:2]
        for _, human in humans.items():
            for part in human:
                if part['partIdx'] not in CocoPairsRender:
                    continue
                center1 = (int((part['c1'][0] + 0.5) * image_w / heat_w), int((part['c1'][1] + 0.5) * image_h / heat_h))
                center2 = (int((part['c2'][0] + 0.5) * image_w / heat_w), int((part['c2'][1] + 0.5) * image_h / heat_h))
                cv2.circle(image, center1, 2, (255, 0, 0), thickness=3, lineType=8, shift=0)
                cv2.circle(image, center2, 2, (255, 0, 0), thickness=3, lineType=8, shift=0)
                image = cv2.line(image, center1, center2, (255, 0, 0), 1)
        cv2.imshow('result', image)
        cv2.waitKey(0)