from __future__ import division
import os
import sys
import time
import csv
import random
from random import randint
from math import floor, radians
from glob import glob
import tensorflow as tf
import tensorflow.contrib.slim as slim
import numpy as np
from six.moves import xrange
from ops import *
from utils import *
from Loss_AMSoftmax import *

SUBJECT_NUM_VGG2 = 8631

class DCGAN(object):
    def __init__(self, sess, config):
        """
        Args:
            sess: TensorFlow session
            batch_size: The size of batch. Should be specified before training.
            output_size: (optional) The resolution in pixels of the images. [64]
            y_dim: (optional) Dimension of dim for y. [None]
            z_dim: (optional) Dimension of dim for Z. [100]
            gf_dim: (optional) Dimension of gen filters in first conv layer. [64]
            df_dim: (optional) Dimension of discrim filters in first conv layer. [64]
            gfc_dim: (optional) Dimension of gen units for for fully connected layer. [1024]
            dfc_dim: (optional) Dimension of discrim units for fully connected layer. [1024]
            c_dim: (optional) Dimension of image color. For grayscale input, set to 1. [3]
        """
        self.sess = sess
        self.is_crop = config.is_crop
        self.is_grayscale = (config.c_dim == 1)
        self.batch_size = config.batch_size
        self.image_size = config.image_size
        self.sample_size = config.sample_size
        self.output_size = config.image_size
        self.gf_dim = config.gf_dim
        self.df_dim = config.df_dim
        self.gfc_dim = config.gfc_dim
        self.dfc_dim = config.dfc_dim
        self.c_dim = config.c_dim

        self.random_flip = True
        self.padding = 0.15
        self.before_crop_size = int(self.output_size * (1 + self.padding))
        self.random_rotate = True

        self.z_dim = config.z_dim
        self.pose_dim = 1
        self.pose_repeat = self.z_dim

        #self.il_dim = il_dim
        # batch normalization : deals with poor initialization helps gradient flow
        self.d_bn0_0 = batch_norm(name='d_k_bn0_0')
        self.d_bn0_1 = batch_norm(name='d_k_bn0_1')
        self.d_bn0_2 = batch_norm(name='d_k_bn0_2')
        self.d_bn1_0 = batch_norm(name='d_k_bn1_0')
        self.d_bn1_1 = batch_norm(name='d_k_bn1_1')
        self.d_bn1_2 = batch_norm(name='d_k_bn1_2')
        self.d_bn1_3 = batch_norm(name='d_k_bn1_3')
        self.d_bn2_0 = batch_norm(name='d_k_bn2_0')
        self.d_bn2_1 = batch_norm(name='d_k_bn2_1')
        self.d_bn2_2 = batch_norm(name='d_k_bn2_2')
        self.d_bn3_0 = batch_norm(name='d_k_bn3_0')
        self.d_bn3_1 = batch_norm(name='d_k_bn3_1')
        self.d_bn3_2 = batch_norm(name='d_k_bn3_2')
        self.d_bn3_3 = batch_norm(name='d_k_bn3_3')
        self.d_bn4_0 = batch_norm(name='d_k_bn4_0')
        self.d_bn4_1 = batch_norm(name='d_k_bn4_1')
        self.d_bn4_2 = batch_norm(name='d_k_bn4_2')
        self.d_bn4_1_l = batch_norm(name='d_k_bn4_1_l')
        self.d_bn4_2_l = batch_norm(name='d_k_bn4_2_l')
        self.d_bn4_1_r = batch_norm(name='d_k_bn4_1_r')
        self.d_bn4_2_r = batch_norm(name='d_k_bn4_2_r')
        self.d_bn4_1_p = batch_norm(name='d_k_bn4_1_pose')
        self.d_bn4_2_p = batch_norm(name='d_k_bn4_2_pose')
        self.d_bn4_1_a = batch_norm(name='d_k_bn4_1_ad')
        self.d_bn4_2_a = batch_norm(name='d_k_bn4_2_ad')
        self.d_bn5   = batch_norm(name='d_k_bn5')

        self.g_bn0_0 = batch_norm(name='g_k_bn0_0')
        self.g_bn0_1 = batch_norm(name='g_k_bn0_1')
        self.g_bn0_2 = batch_norm(name='g_k_bn0_2')
        self.g_bn0_3 = batch_norm(name='g_k_bn0_3')
        self.g_bn1_0 = batch_norm(name='g_k_bn1_0')
        self.g_bn1_1 = batch_norm(name='g_k_bn1_1')
        self.g_bn1_2 = batch_norm(name='g_k_bn1_2')
        self.g_bn1_3 = batch_norm(name='g_k_bn1_3')
        self.g_bn2_0 = batch_norm(name='g_k_bn2_0')
        self.g_bn2_1 = batch_norm(name='g_k_bn2_1')
        self.g_bn2_2 = batch_norm(name='g_k_bn2_2')
        self.g_bn2_3 = batch_norm(name='g_k_bn2_3')
        self.g_bn3_0 = batch_norm(name='g_k_bn3_0')
        self.g_bn3_1 = batch_norm(name='g_k_bn3_1')
        self.g_bn3_2 = batch_norm(name='g_k_bn3_2')
        self.g_bn3_3 = batch_norm(name='g_k_bn3_3')
        self.g_bn4_0 = batch_norm(name='g_k_bn4_0')
        self.g_bn4_1 = batch_norm(name='g_k_bn4_1')
        self.g_bn4_2 = batch_norm(name='g_k_bn4_2')
        self.g_bn4_c = batch_norm(name='g_h_bn4_c')
        self.g_bn5   = batch_norm(name='g_k_bn5')

        self.g1_bn0_0 = batch_norm(name='g_h_bn0_0')
        self.g1_bn0_1 = batch_norm(name='g_h_bn0_1')
        self.g1_bn0_2 = batch_norm(name='g_h_bn0_2')
        self.g1_bn1_0 = batch_norm(name='g_h_bn1_0')
        self.g1_bn1_1 = batch_norm(name='g_h_bn1_1')
        self.g1_bn1_2 = batch_norm(name='g_h_bn1_2')
        self.g1_bn2_0 = batch_norm(name='g_h_bn2_0')
        self.g1_bn2_1 = batch_norm(name='g_h_bn2_1')
        self.g1_bn2_2 = batch_norm(name='g_h_bn2_2')
        self.g1_bn3_0 = batch_norm(name='g_h_bn3_0')
        self.g1_bn3_1 = batch_norm(name='g_h_bn3_1')
        self.g1_bn3_2 = batch_norm(name='g_h_bn3_2')
        self.g1_bn4_0 = batch_norm(name='g_h_bn4_0')
        self.g1_bn4   = batch_norm(name='g_h_bn4')
        self.g1_bn5   = batch_norm(name='g_h_bn5')
        self.dataset_name = config.dataset
        self.checkpoint_dir = config.checkpoint_dir
        self.samples_dir = config.samples_dir
        model_dir = "%s_%s_%s_%s_%s_%s_%s" % (self.dataset_name, self.batch_size, self.output_size, self.gf_dim, self.gfc_dim, self.df_dim, self.dfc_dim)
        if not os.path.exists(self.samples_dir+"/"+model_dir):
            os.makedirs(self.samples_dir+"/"+model_dir)
        if not os.path.exists(self.checkpoint_dir+"/"+model_dir):
            os.makedirs(self.checkpoint_dir+"/"+model_dir)
        self.build_model()

    def build_model(self):

        self.subject_num = SUBJECT_NUM_VGG2
        self.multi_N = 4

        self.input_labels= tf.placeholder(tf.int64, [self.batch_size,], name='positive_labels')
        self.input_filenames = [tf.placeholder(dtype=tf.string) for _ in range(self.batch_size)]

        self.real_positive_labels= tf.placeholder(tf.int64, [self.batch_size,], name='positive_hot_code_labels')
        self.real_negative_labels= tf.placeholder(tf.int64, [self.batch_size,], name='negative_hot_code_labels')
        self.real_anchor_hot_code_poses= tf.placeholder(tf.float32, [self.batch_size, self.pose_dim], name='anchor_hot_code_poses')
        #self.real_positive_hot_code_poses= tf.placeholder(tf.float32, [self.batch_size, self.pose_dim], name='positive_hot_code_poses')
        #self.real_positive2_hot_code_poses= tf.placeholder(tf.float32, [self.batch_size, self.pose_dim], name='positive2_hot_code_poses')
        #self.real_positive3_hot_code_poses= tf.placeholder(tf.float32, [self.batch_size, self.pose_dim], name='positive3_hot_code_poses')
        #self.real_negative_hot_code_poses= tf.placeholder(tf.float32, [self.batch_size, self.pose_dim], name='negative_hot_code_poses')
        self.real_aflw_hot_code_poses= tf.placeholder(tf.float32, [self.batch_size, self.pose_dim], name='anchor_hot_code_poses')

        self.aflw_input_filenames = [tf.placeholder(dtype=tf.string) for _ in range(self.batch_size)]
        self.anchor_input_filenames = [tf.placeholder(dtype=tf.string) for _ in range(self.batch_size)]
        self.positive_input_filenames = [tf.placeholder(dtype=tf.string) for _ in range(self.batch_size)]
        self.positive2_input_filenames = [tf.placeholder(dtype=tf.string) for _ in range(self.batch_size)]
        self.positive3_input_filenames = [tf.placeholder(dtype=tf.string) for _ in range(self.batch_size)]
        self.negative_input_filenames = [tf.placeholder(dtype=tf.string) for _ in range(self.batch_size)]

        self.G_anchor_hot_code_poses= tf.placeholder(tf.float32, [self.batch_size, self.pose_dim], name='G_anchor_hot_code_poses')
        self.G_positive_hot_code_poses= tf.placeholder(tf.float32, [self.batch_size, self.pose_dim], name='G_positive_hot_code_poses')
        self.G_positive2_hot_code_poses= tf.placeholder(tf.float32, [self.batch_size, self.pose_dim], name='G_positive2_hot_code_poses')
        self.G_positive3_hot_code_poses= tf.placeholder(tf.float32, [self.batch_size, self.pose_dim], name='G_positive3_hot_code_poses')
        self.G_negative_hot_code_poses= tf.placeholder(tf.float32, [self.batch_size, self.pose_dim], name='G_negative_hot_code_poses')
        self.G_combine_hot_code_poses= tf.placeholder(tf.float32, [self.batch_size, self.pose_dim], name='G_combine_hot_code_poses')

        #self.sample_images= tf.placeholder(tf.float32, [self.sample_size] + [self.output_size, self.output_size, self.c_dim], name='sample_images')
        #self.z = tf.placeholder(tf.float32, [None, self.z_dim], name='z')
        #self.pose = tf.placeholder(tf.float32, [None, self.pose_dim], name='pose')
        #self.il = tf.placeholder(tf.float32, [None, self.il_dim], name='il')
        #self.sample_input_images = tf.placeholder(tf.float32, [1, self.output_size, self.output_size, self.c_dim], name='sample_input_images')
        #self.k52 = tf.placeholder(tf.float32, [1, self.gfc_dim], name='sample_k52')

        # Networks
        def filename2image(input_filenames, batch_size, isflip = True):
            images = []
            for i in range(batch_size):
                file_contents = tf.read_file(input_filenames[i])
                image = tf.image.decode_jpeg(file_contents, channels=3)
                image = tf.image.resize_images(image, [self.before_crop_size, self.before_crop_size])
                if self.random_rotate:
                    image = tf.py_func(random_rotate_image, [image], tf.uint8)
                if (self.padding > 0):
                    image = tf.random_crop(image, [self.output_size, self.output_size, 3])
                #if args.random_rotate:
                #image = tf.py_func(facenet.random_rotate_image, [image], tf.uint8)
                #if args.random_crop:
                #    image = tf.random_crop(image, [args.image_size, args.image_size, 3])
                #else:
                #    image = tf.image.resize_image_with_crop_or_pad(image, args.image_size, args.image_size)
                if self.random_flip and isflip:
                    image = tf.image.random_flip_left_right(image)
                images.append(tf.subtract(tf.div(tf.cast(image, dtype=tf.float32), 127.5), 1.0))
            return tf.stack(images)

        self.aflw_input_images      = filename2image(self.aflw_input_filenames, self.batch_size, isflip=False)
        self.anchor_input_images    = filename2image(self.anchor_input_filenames, self.batch_size)
        self.positive_input_images  = filename2image(self.positive_input_filenames, self.batch_size)
        self.positive2_input_images = filename2image(self.positive2_input_filenames, self.batch_size)
        self.positive3_input_images = filename2image(self.positive3_input_filenames, self.batch_size)
        self.negative_input_images  = filename2image(self.negative_input_filenames, self.batch_size)

        self.G_encode_anchor, self.G_coefficient_anchor = self.generator_encoder( self.anchor_input_images, is_reuse=False)
        self.G_decode_anchor = self.generator_decoder( tf.random_uniform( shape=[self.batch_size, self.z_dim], minval=-1, maxval=1),
                                                       self.G_encode_anchor,  self.G_anchor_hot_code_poses,  is_reuse=False)

        self.G_encode_positive, self.G_coefficient_positive= self.generator_encoder( self.positive_input_images, is_reuse=True)
        self.G_encode_positive2, self.G_coefficient_positive2 = self.generator_encoder( self.positive2_input_images, is_reuse=True)
        self.G_encode_positive3, self.G_coefficient_positive3 = self.generator_encoder( self.positive3_input_images, is_reuse=True)
        self.G_encode_negative, _                             = self.generator_encoder( self.negative_input_images, is_reuse=True)

        self.G_encode_combine = tf.div( tf.multiply(self.G_encode_anchor, self.G_coefficient_anchor)       + tf.multiply(self.G_encode_positive, self.G_coefficient_positive) + \
                                        tf.multiply(self.G_encode_positive2, self.G_coefficient_positive2) + tf.multiply(self.G_encode_positive3, self.G_coefficient_positive3) ,
                                        self.G_coefficient_anchor + self.G_coefficient_positive + self.G_coefficient_positive2 + self.G_coefficient_positive3 + 1e-6 )

        self.G_decode_positive = self.generator_decoder( tf.random_uniform( shape=[self.batch_size, self.z_dim], minval=-1, maxval=1), self.G_encode_positive, self.G_positive_hot_code_poses, is_reuse=True)
        self.G_decode_positive2 = self.generator_decoder( tf.random_uniform( shape=[self.batch_size, self.z_dim], minval=-1, maxval=1), self.G_encode_positive2, self.G_positive2_hot_code_poses, is_reuse=True)
        self.G_decode_positive3 = self.generator_decoder( tf.random_uniform( shape=[self.batch_size, self.z_dim], minval=-1, maxval=1), self.G_encode_positive3, self.G_positive3_hot_code_poses, is_reuse=True)
        self.G_decode_negative = self.generator_decoder( tf.random_uniform( shape=[self.batch_size, self.z_dim], minval=-1, maxval=1), self.G_encode_negative, self.G_negative_hot_code_poses, is_reuse=True)
        self.G_decode_combine  = self.generator_decoder( tf.random_uniform( shape=[self.batch_size, self.z_dim], minval=-1, maxval=1), self.G_encode_combine,  self.G_combine_hot_code_poses,  is_reuse=True)

        self.D_all_logits, self.D_all_pose_logits, self.D_all_id_logits, self.D_all_k5_ids, self.D_all_k6_ws = self.discriminator(tf.concat(axis=0,
            values=[self.aflw_input_images,
                    self.anchor_input_images,
                    self.positive_input_images,
                    self.positive2_input_images,
                    self.positive3_input_images,
                    self.negative_input_images,
                    self.G_decode_anchor,
                    self.G_decode_positive,
                    self.G_decode_positive2,
                    self.G_decode_positive3,
                    self.G_decode_combine,
                    self.G_decode_negative]), is_reuse=False)

        self.D_R_logits, self.D_G_logits = tf.split(axis=0, num_or_size_splits=2, value=self.D_all_logits)
        self.D_R_id_logits, self.D_G_id_logits = tf.split(axis=0, num_or_size_splits=2, value=self.D_all_id_logits)
        self.D_R_pose_logits, self.D_G_pose_logits = tf.split(axis=0, num_or_size_splits=2, value=self.D_all_pose_logits)

        self.D_R_pose_awlf_logits = tf.slice(self.D_R_pose_logits, [0,0], [self.batch_size,-1])
        self.D_R_id_casia_logits  = tf.slice(self.D_R_id_logits, [self.batch_size,0], [self.batch_size * 4,-1])
        self.d_acc = slim.metrics.accuracy(tf.argmax(self.D_R_id_casia_logits, 1),
                                           tf.concat(axis=0, values=[ self.real_positive_labels,
                                                                      self.real_positive_labels,
                                                                      self.real_positive_labels,
                                                                      self.real_positive_labels]), weights=100.0)
        self.D_R_id_casia_logits = AMSoftmax_logit_v2(
            self.D_R_id_casia_logits,
            self.D_all_k6_ws,
            label_batch=tf.concat(axis=0, values=[self.real_positive_labels,
                    self.real_positive_labels,
                    self.real_positive_labels,
                    self.real_positive_labels]), nrof_classes=self.subject_num)
        self.D_G_id_logits = tf.slice(self.D_G_id_logits, [0, 0], [self.batch_size * 5, -1])
        self.g_acc = slim.metrics.accuracy(tf.argmax(self.D_G_id_logits, 1), tf.concat(axis=0,
                                                                                values=[
                                                                                self.real_positive_labels,
                                                                                self.real_positive_labels,
                                                                                self.real_positive_labels,
                                                                                self.real_positive_labels,
                                                                                self.real_positive_labels]), weights=100.0)
        self.D_G_id_logits = AMSoftmax_logit_v2(
            self.D_G_id_logits,
            self.D_all_k6_ws,
            label_batch=tf.concat(axis=0, values=[self.real_positive_labels,
                                                  self.real_positive_labels,
                                                  self.real_positive_labels,
                                                  self.real_positive_labels,
                                                  self.real_positive_labels]), nrof_classes=self.subject_num)

        #self.samplers, _, _ = self.sampler(tf.random_uniform( shape=[self.batch_size, self.z_dim], minval=-1, maxval=1), self.sample_images, self.pose)
        #self.test_sampler, self.sampler_feature, self.sampler_coefficient = self.sampler(tf.random_uniform( shape=[self.batch_size, self.z_dim], minval=-1, maxval=1), self.sample_input_images, self.pose)

        #self.k5, _ = self.generator_encoder(self.sample_input_images, is_reuse=True, is_training = False)
        #self.h0 = self.generator_decoder(self.z, self.k52, self.pose, is_reuse=True, is_training=False)
        #self.k5, _ = self.generator_encoder(self.h0, is_reuse=True, is_training = False)
        #_,_,_,self.k5 = self.discriminator(self.h0, is_reuse=True, is_training = False)

        #self.D_R_logits3 = tf.concat(0,[self.D_R_logits, D_R_logits2])
        #self.D_R_logits3 = tf.concat(0,[self.D_R_logits, D_R_logits2])

        # D Loss
        self.d_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=self.D_R_logits, labels=0.999*tf.ones_like(self.D_R_logits)))
        self.d_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=self.D_G_logits, labels=tf.zeros_like(self.D_G_logits)))
        #self.d_loss_real_pose = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=self.D_R_pose_awlf_logits, labels=self.real_aflw_hot_code_poses))
        self.d_loss_real_pose = tf.nn.l2_loss(self.D_R_pose_awlf_logits - self.real_aflw_hot_code_poses) / self.pose_dim / self.batch_size

        self.d_loss_real_id = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=self.D_R_id_casia_logits, labels=tf.concat(axis=0,
            values=[self.real_positive_labels,
                    self.real_positive_labels,
                    self.real_positive_labels,
                    self.real_positive_labels])))

        self.d_loss =  self.d_loss_real + self.d_loss_fake + self.d_loss_real_id + self.d_loss_real_pose

        # G Loss
        self.g_loss_d_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=self.D_G_logits, labels=tf.ones_like(self.D_G_logits)))
        self.g_loss_d_id   = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=self.D_G_id_logits, labels=tf.concat(axis=0,
            values=[self.real_positive_labels,
                    self.real_positive_labels,
                    self.real_positive_labels,
                    self.real_positive_labels,
                    self.real_positive_labels])))
        self.g_loss_d_id_2 = tf.zeros(1)#tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(tf.concat(0,[self.G_anchor_id_logits, self.G_positive_id_logits, self.G_positive2_id_logits, self.G_positive3_id_logits]), tf.concat(0,[self.real_positive_hot_code_labels, self.real_positive_hot_code_labels, self.real_positive_hot_code_labels, self.real_positive_hot_code_labels])))
        self.g_loss_d_pose = tf.nn.l2_loss(self.D_G_pose_logits - tf.concat(axis=0, values=[self.G_anchor_hot_code_poses, self.G_positive_hot_code_poses, self.G_positive2_hot_code_poses, self.G_positive3_hot_code_poses, self.G_combine_hot_code_poses, self.G_negative_hot_code_poses]))  / self.pose_dim / self.batch_size / 6

        self.g_loss =self.g_loss_d_real + self.g_loss_d_id +  self.g_loss_d_pose #+   + 2*self.g_loss_d_id_2

        self.g_acc2 = tf.zeros(1)

        # Sumaries

        d_loss_sum = tf.summary.scalar("D loss", self.d_loss)
        d_loss_id_sum = tf.summary.scalar("D_id", self.d_loss_real_id)
        d_acc_id_sum = tf.summary.scalar("D_id_acc", self.d_acc)
        d_loss_pose_sum = tf.summary.scalar("D_pose", self.d_loss_real_pose)
        d_loss_gan_r_sum = tf.summary.scalar("D_gan_r", self.d_loss_real)
        d_loss_gan_f_sum = tf.summary.scalar("D_gan_f", self.d_loss_fake)

        g_loss_sum = tf.summary.scalar("G loss", self.g_loss)
        g_loss_id_sum = tf.summary.scalar("G_id", self.g_loss_d_id)
        g_loss_pose_sum = tf.summary.scalar("G_pose", self.g_loss_d_pose)
        g_loss_gan_sum = tf.summary.scalar("G_gan", self.g_loss_d_real)
        g_acc_id_sum = tf.summary.scalar("G_id_acc", self.g_acc)

        d_real_sum = tf.summary.histogram("d_R", tf.sigmoid(self.D_R_logits))
        d_fake_sum = tf.summary.histogram("d_G", tf.sigmoid(self.D_G_logits))
        g_out_img_summary = tf.summary.image("Generated images", self.G_decode_anchor, max_outputs=5)
        g_in_img_summary = tf.summary.image("Input", self.anchor_input_images, max_outputs=5)

        self.d_summary_op = tf.summary.merge([d_loss_sum, d_loss_id_sum, d_acc_id_sum, d_loss_pose_sum, d_loss_gan_r_sum, d_loss_gan_f_sum, d_real_sum])
        self.g_summary_op = tf.summary.merge([g_loss_sum, g_loss_id_sum, g_acc_id_sum, g_loss_pose_sum, g_loss_gan_sum, d_fake_sum])
        self.g_img_summary_op = tf.summary.merge([g_in_img_summary, g_out_img_summary])
        self.summary_writer = tf.summary.FileWriter(self.checkpoint_dir+"/"+self.model_dir+"/log", self.sess.graph)

        #self.d_loss_real_sum = tf.scalar_summary("d_loss_real", self.d_loss_real)
        #self.d_loss_fake_sum = tf.scalar_summary("d_loss_fake", self.d_loss_fake)

        #self.d_loss = self.d_loss_real + self.d_loss_fake
        #self.g_loss_sum = tf.scalar_summary("g_loss", self.g_loss)
        #self.d_loss_sum = tf.scalar_summary("d_loss", self.d_loss)

        t_vars = tf.trainable_variables()
        self.d_vars = [var for var in t_vars if 'd_' in var.name]
        self.d_vars_load = [var for var in tf.global_variables() if 'd_' in var.name]
        self.g_vars = [var for var in t_vars if 'g_' in var.name]
        self.g_en_vars = [var for var in tf.global_variables() if 'g_k' in var.name]
        self.g_de_vars = [var for var in tf.global_variables() if 'g_h' in var.name]

        self.saver = tf.train.Saver(keep_checkpoint_every_n_hours=1, max_to_keep = 20)
        self.d_saver = tf.train.Saver([var for var in self.d_vars_load if not ('pose' in var.name or 'ad' in var.name)], keep_checkpoint_every_n_hours=1, max_to_keep = 0)
        #self.d_saver = tf.train.Saver([var for var in self.d_vars] , keep_checkpoint_every_n_hours=1, max_to_keep = 0)

        def name_in_checkpoint(var):
            print(var.op.name)
            print('d' + var.op.name[1:])
            return 'd' + var.op.name[1:]

        self.variables_to_restore = {name_in_checkpoint(var):var for var in self.g_en_vars}
        self.g_en_saver = tf.train.Saver(self.variables_to_restore, keep_checkpoint_every_n_hours=1, max_to_keep = 0)
        #self.g_de_saver = tf.train.Saver(self.g_de_vars, keep_checkpoint_every_n_hours=1, max_to_keep = 0)
        #self.g_saver = tf.train.Saver([var for var in self.g_vars if not 'g_h5_lin' in var.name], keep_checkpoint_every_n_hours=1, max_to_keep = 0)

    def test_IJBA(self, config, isReload = False, iteration = 0, isFlip = True):

        if isReload:
            could_load, iteration = self.load(self.checkpoint_dir)
            if could_load:
                print(" [*] Load SUCCESS")
            else:
                print(" [!] Load failed...")

        if  config.dataset == 'CASIA' or config.dataset == 'CASIA_MTPIE':
            # Single image IJBA
            images = load_IJBA_recrop_test();

            N = images.shape[0]
            out_features = np.zeros((N, self.gfc_dim), dtype=np.float32)
            out_coefficients = np.zeros((N, 1), dtype=np.float32)

            target_pose = np.zeros([1, self.pose_dim], dtype=np.float)
            #target_pose[0,self.pose_dim/2] = 1.0
            sample_z = np.zeros(shape=(1 , self.z_dim))

            gallery_features = np.zeros((N, self.gfc_dim), dtype=np.float32)
            gallery_coefficients = np.zeros((N, 1), dtype=np.float32)
            #gallery_rotated_images = np.zeros((2, 96, 96, 3), dtype=np.float32)

            idxes = range(0,N)

            bar = Bar('Processing', max=len(idxes), suffix='%(percent)d%%')

            co = 0

            idx = 0

            for co in range(0,N):
                idx = idxes[co]
                ##co = co + 1
                if (co > len(idxes)):
                    idx = input("Input");

                batch_img = np.array(images[idx,:,:,:]).astype(np.float32).reshape((1,self.output_size, self.output_size, self.c_dim))
                batch_img = np.array(batch_img).reshape(1,96,96,3)

                gallery_rotated_image, gallery_feature, gallery_coefficient = self.sess.run([ self.test_sampler, self.sampler_feature, self.sampler_coefficient], feed_dict={self.z: sample_z, self.sample_input_images:batch_img, self.pose: target_pose} )

                if isFlip:
                    batch_img = np.flip(batch_img, axis = 2)
                    gallery_feature2, gallery_coefficient2 = self.sess.run([ self.sampler_feature, self.sampler_coefficient], feed_dict={self.z: sample_z, self.sample_input_images:batch_img, self.pose: target_pose} )
                    gallery_feature = (gallery_feature + gallery_feature2)/2.0
                    gallery_coefficient = (gallery_coefficient + gallery_coefficient2)/2.0

                gallery_features[idx,:] = gallery_feature
                gallery_coefficients[idx,:] = gallery_coefficient

                    #gallery_features2[i,:] = gallery_feature2
                    #gallery_rotated_images[i,:,:,:] =  np.array(gallery_rotated_image).reshape(1, self.output_size, self.output_size, self.c_dim)
                bar.next()
                #gallery_rotated_images[1,:,:,:] = batch_img
                #gallery_rotated_images[-1,:,:,:] = batch_img2
                #gallery_rotated_images[2,:,:,:] = np.array(images[frontal_idx[idx-5000]-1,:,:,:]).astype(np.float32).reshape((1,self.output_size, self.output_size, self.c_dim))

                #save_images(gallery_rotated_images, [1, gallery_rotated_images.shape[0]],'IJBA_single/rotated_images_{:d}.png'.format(idx))# .format(idx))
                #gallery_feature.tofile('IJBA_single/feature_{:d}.bin'.format(idx))# .format(idx))

            checkpoint_dir = os.path.join(self.checkpoint_dir, self.model_dir)
            if isFlip:
                np.savetxt("%s/IJBA_features_iter_%d_flip.txt" % (checkpoint_dir, iteration), gallery_features)
                np.savetxt("%s/IJBA_coefficient_iter_%d_flip.txt" % (checkpoint_dir, iteration), gallery_coefficients)

                print("\nWrited to %s/IJBA_features_iter_%d_flip.txt" % (checkpoint_dir, iteration))
            else:
                np.savetxt("%s/IJBA_features_iter_%d.txt" % (checkpoint_dir, iteration), gallery_features)
                np.savetxt("%s/IJBA_coefficient_iter_%d.txt" % (checkpoint_dir, iteration), gallery_coefficients)
                print("\nWrited to %s/IJBA_features_iter_%d.txt" % (checkpoint_dir, iteration))

    def cos_loop(self,matrix, vector):
        """
        Calculating pairwise cosine distance using a common for loop with manually calculated cosine value.
        """
        neighbors = []
        for row in range(matrix.shape[0]):
            vector_norm = np.linalg.norm(vector)
            row_norm = np.linalg.norm(matrix[row,:])
            cos_val = vector.dot(matrix[row,:]) / (vector_norm * row_norm)
            neighbors.append(cos_val)
        return neighbors

    def train(self, config):

        #tf.set_random_seed(666)
        #np.random.seed(777)

        if config.dataset == 'VGG2':
            path_VGG2_l, pid_VGG2_l, path_VGG2_n, pid_VGG2_n = load_VGG2_train_by_list(initial_id_l=0, initial_id_n=0)

            path_all_l = path_VGG2_l
            path_all = path_VGG2_n

            pid_all = pid_VGG2_n

            path_AFLW, roll_AFLW, pitch_AFLW, yaw_AFLW = load_AFLW_by_list()

            num_VGG2 = len(pid_all)
            num_AFLW = len(yaw_AFLW)

            id_dict = {}
            for i in range(num_VGG2):
                if (pid_all[i] not in id_dict):
                    id_dict[pid_all[i]] = []
                id_dict[pid_all[i]].append(i)

            hot_code_pose_AFLW = np.zeros((num_AFLW, self.pose_dim), dtype=np.float)
            for i, label in enumerate(yaw_AFLW):
                #pose_code = int(floor((yaw_AFLW[i]+97.5)/15))
                #if (pose_code < 0):
                #    pose_code = 0
                #elif (pose_code >= self.pose_dim):
                #    pose_code = self.pose_dim-1
                #hot_code_pose_AFLW[i, pose_code] = 1.0
                hot_code_pose_AFLW[i,0] = radians(yaw_AFLW[i])

        else:
            data = glob(os.path.join("./data", config.dataset, "*.jpg"))
        #np.random.shuffle(data)
        d_optim = tf.train.AdamOptimizer(config.learning_rate, beta1=config.beta1) \
                          .minimize(self.d_loss, var_list=self.d_vars)
        g_optim = tf.train.AdamOptimizer(config.learning_rate, beta1=config.beta1) \
                          .minimize(self.g_loss, var_list=self.g_vars)
        tf.global_variables_initializer().run()

        #self.g_sum = tf.merge_summary([self.z_sum, self.d__sum,
        #    self.G_sum, self.d_loss_fake_sum, self.g_loss_sum])
        #self.d_sum = tf.merge_summary([self.z_sum, self.d_sum, self.d_loss_real_sum, self.d_loss_sum])
        #self.writer = tf.train.SummaryWriter("./logs", self.sess.graph)
        sample_z = np.random.uniform(-1, 1, size=(self.sample_size , self.z_dim))

        """ Sample images """
        '''
        if config.dataset == 'mnist':
            sample_images = data_X[0:self.sample_size]
            sample_labels = data_y[0:self.sample_size]
        elif config.dataset == 'CASIA' or config.dataset == 'CASIA_MTPIE':
            STEP = 100
            sample_files = range(0, self.sample_size)
            for i in range(0, self.sample_size):
                sample_files[i] = int((i/self.pose_dim)*STEP+i)

            sample = [random_crop(images_CASIA[sample_file,:,:], self.output_size, self.output_size) for sample_file in sample_files]
            sample_images = np.array(sample).astype(np.float32)/127.5-1

            sample_pose = np.zeros((self.sample_size, self.pose_dim), dtype=np.float)
            for i in range(self.sample_size):
                sample_pose[i] = radians( i - 80.0)

            ##save_images(sample_images, [8, 8], './{:s}/{:s}/batch_sample.png'.format(self.sample_dir, self.model_dir))
        elif config.dataset == 'MTPIE':
            sample_labels, sample_images, sample_images_wrong = load_MTPIE_samples();
            sample_pose = np.zeros((64, self.pose_dim), dtype=np.float)
            for i in range(1,self.sample_size):
                sample_images[i,:,:,:] = sample_images[i / self.pose_dim,:,:,:]
            for i in range(self.sample_size):
                sample_pose[i,i % self.pose_dim] = 1.0
            sample_il = np.zeros((self.sample_size, self.il_dim), dtype=np.float)
            for i in range(self.sample_size):
                sample_il[i,i % selfil_dim] = 1.0
        else:
            sample_files = data[0:self.sample_size]
            sample = [get_image(sample_file, self.image_size, is_crop=self.is_crop, resize_w=self.output_size, is_grayscale = self.is_grayscale) for sample_file in sample_files]
            if (self.is_grayscale):
                sample_images = np.array(sample).astype(np.float32)[:, :, :, None]
            else:
                sample_images = np.array(sample).astype(np.float32)
        '''

        self.d_saver = tf.train.Saver([var for var in self.d_vars_load if not ('pose' in var.name or 'ad' in var.name)], keep_checkpoint_every_n_hours=1, max_to_keep=0)
        self.g_en_saver = tf.train.Saver(self.variables_to_restore, keep_checkpoint_every_n_hours=1, max_to_keep=0)
        #self.g_de_saver = tf.train.Saver(self.g_de_vars, keep_checkpoint_every_n_hours=1, max_to_keep=0)
        self.saver = tf.train.Saver(keep_checkpoint_every_n_hours=1, max_to_keep=20)

        """Train DCGAN"""
        could_load, checkpoint_counter = self.load(self.checkpoint_dir)
        if could_load:
            counter = checkpoint_counter
            print(" [*] Load SUCCESS")
        else:
            counter = 0
            print(" [!] Load failed...")
        #self.d_saver.restore(self.sess, "/media/tranluan/CropFaceData/Repos/DR-GAN_Bangjie_0524/checkpoints/DRGAN_n4/0525/DCGAN.model-789000")
        #self.g_de_saver.restore(self.sess, "/media/tranluan/CropFaceData/Repos/DR-GAN_Bangjie_0524/checkpoints/DRGAN_n4/0525/DCGAN.model-789000")

        start_time = time.time()

        for epoch in xrange(config.epoch):
            valid_idx = np.random.permutation(num_VGG2)
            batch_idxs = min(len(valid_idx), config.train_size) // (config.batch_size *  self.multi_N)

            for idx in xrange(0, batch_idxs):
                if config.dataset == 'VGG2':
                    ffeed_dict={}

                    #Random pose code
                    batch_random_anchor_hot_code_poses    = np.random.normal(loc=0, scale=radians(45.0), size=[config.batch_size, self.pose_dim]) #np.random.uniform(low=radians(-90.0), high=radians(90.0), size=[config.batch_size, self.pose_dim])
                    batch_random_positive_hot_code_poses  = np.random.uniform(low=radians(-90.0), high=radians(90.0), size=[config.batch_size, self.pose_dim])
                    batch_random_positive2_hot_code_poses = np.random.uniform(low=radians(-90.0), high=radians(90.0), size=[config.batch_size, self.pose_dim])
                    batch_random_positive3_hot_code_poses = np.random.uniform(low=radians(-90.0), high=radians(90.0), size=[config.batch_size, self.pose_dim])
                    batch_random_negative_hot_code_poses  = np.random.uniform(low=radians(-90.0), high=radians(90.0), size=[config.batch_size, self.pose_dim])
                    batch_random_combine_hot_code_poses   = np.random.uniform(low=radians(-90.0), high=radians(90.0), size=[config.batch_size, self.pose_dim])

                    #batch_random_combine_hot_code_poses   = np.random.uniform(low=radians(-90.0), high=radians(90.0), size=[config.batch_size, self.pose_dim])

                    ffeed_dict[self.G_anchor_hot_code_poses]    = batch_random_anchor_hot_code_poses
                    ffeed_dict[self.G_positive_hot_code_poses]  = batch_random_positive_hot_code_poses
                    ffeed_dict[self.G_positive2_hot_code_poses] = batch_random_positive2_hot_code_poses
                    ffeed_dict[self.G_positive3_hot_code_poses] = batch_random_positive3_hot_code_poses
                    ffeed_dict[self.G_negative_hot_code_poses]  = batch_random_negative_hot_code_poses
                    ffeed_dict[self.G_combine_hot_code_poses]   = batch_random_combine_hot_code_poses

                    #AFLW
                    batch = []
                    aflw_idx = np.random.random_integers(0, num_AFLW-1, self.batch_size)
                    batch_aflw_images = [path_AFLW[batch_file]  for batch_file in aflw_idx]
                    batch_aflw_hot_code_poses    = [hot_code_pose_AFLW[batch_file,:] for batch_file in aflw_idx]
                    ffeed_dict[self.real_aflw_hot_code_poses] = batch_aflw_hot_code_poses

                    for i in range(self.batch_size):
                        ffeed_dict[self.aflw_input_filenames[i]] = batch_aflw_images[i]

                    #VGG2
                    anchor_idx = [0] * config.batch_size
                    positive_idx = [0] * config.batch_size
                    positive2_idx = [0] * config.batch_size
                    positive3_idx = [0] * config.batch_size
                    negative_idx = [0] * config.batch_size

                    ids = id_dict.keys()
                    random.shuffle(ids)
                    for i in range(config.batch_size):
                        idx_list = id_dict.get(ids[i])
                        anchor_idx[i]   = idx_list[random.randint(0, len(idx_list)-1)]
                        positive_idx[i] = idx_list[random.randint(0, len(idx_list)-1)]
                        positive2_idx[i] = idx_list[random.randint(0, len(idx_list)-1)]
                        positive3_idx[i] = idx_list[random.randint(0, len(idx_list)-1)]

                    for i in range(config.batch_size):
                        negative_idx[i] = random.randint(0, len(path_all_l)-1)

                    batch_positive_labels  = [pid_all[batch_file] for batch_file in positive_idx]
                    batch_negative_labels  = np.zeros(config.batch_size)

                    batch_anchor_images    = [path_all[batch_file]  for batch_file in anchor_idx]
                    batch_positive_images  = [path_all[batch_file]  for batch_file in positive_idx]
                    batch_positive2_images = [path_all[batch_file]  for batch_file in positive2_idx]
                    batch_positive3_images = [path_all[batch_file]  for batch_file in positive3_idx]
                    #batch_positive4_images = [path_all[batch_file]  for batch_file in positive4_idx]
                    batch_negative_images  = [path_all_l[batch_file]  for batch_file in negative_idx]

                    ffeed_dict[self.real_positive_labels]       = batch_positive_labels
                    ffeed_dict[self.real_negative_labels]       = batch_negative_labels

                    for i in range(self.batch_size):
                        ffeed_dict[   self.anchor_input_filenames[i]] = batch_anchor_images[i]
                        ffeed_dict[ self.positive_input_filenames[i]] = batch_positive_images[i]
                        ffeed_dict[self.positive2_input_filenames[i]] = batch_positive2_images[i]
                        ffeed_dict[self.positive3_input_filenames[i]] = batch_positive3_images[i]
                        ffeed_dict[ self.negative_input_filenames[i]] = batch_negative_images[i]

                    #batch_anchor_hot_code_poses    = [hot_code_pose[batch_file,:] for batch_file in anchor_idx]
                    #batch_positive_hot_code_poses  = [hot_code_pose[batch_file,:] for batch_file in positive_idx]
                    #batch_positive2_hot_code_poses  = [hot_code_pose[batch_file,:] for batch_file in positive2_idx]
                    #batch_positive3_hot_code_poses  = [hot_code_pose[batch_file,:] for batch_file in positive3_idx]
                    #batch_negative_hot_code_poses  = [hot_code_pose[batch_file,:] for batch_file in negative_idx]

                try:

                    if counter % (2) == 0:
                        # Update D network
                        _, summary_str = self.sess.run([d_optim,  self.d_summary_op], feed_dict= ffeed_dict)
                        self.summary_writer.add_summary(summary_str, global_step=counter)

                        self.sess.run([g_optim], feed_dict= ffeed_dict)

                    else:
                        # Update G network
                        _, summary_str = self.sess.run([g_optim, self.g_summary_op], feed_dict= ffeed_dict)
                        self.summary_writer.add_summary(summary_str, global_step=counter)
                except:
                    #print("Unexpected error:" + sys.exc_info()[0])
                    continue

                counter += 1
                if np.mod(counter, 10) == 1:
                    _, errD_real_pose, errD_real_id, errD_real, errD_fake, errD, errG_d_id, errG_d_id2, errG_d_pose, errG_d_id, errG_d_real, errG, d_acc, g_acc  = \
                    self.sess.run([g_optim, self.d_loss_real_pose, self.d_loss_real_id, self.d_loss_real, self.d_loss_fake, self.d_loss, self.g_loss_d_id, self.g_loss_d_id_2, self.g_loss_d_pose, self.g_loss_d_id, self.g_loss_d_real, self.g_loss, self.d_acc, self.g_acc], feed_dict= ffeed_dict)

                    print("Epoch: [%2d] [%4d/%4d] time: %4.1f, d_loss: %.4f (pose:%.3f, id: %.3f (R: %.2f), err_G: %.3f (real: %.2f, fake: %.2f) ), g_loss: %.4f (real:%.3f, id:%.3f (R: %.2f), pose:%.3f)" \
                       %(epoch, idx, batch_idxs, time.time() - start_time, errD, errD_real_pose, errD_real_id, d_acc, errD_real + errD_fake, errD_real, errD_fake, errG, errG_d_real, errG_d_id, g_acc, errG_d_pose ))

                if np.mod(counter, 500) == 1:
                    summary_str = self.sess.run(self.g_img_summary_op, feed_dict= ffeed_dict)
                    self.summary_writer.add_summary(summary_str, global_step=counter)

                '''
                if np.mod(counter, 500) == 1:

                    samples = self.sess.run( self.samplers, feed_dict={self.z: sample_z, self.sample_images: sample_images, self.pose: sample_pose} )
                    if (counter == 1):
                        save_images(sample_images, [13, 13],'./{:s}/{:s}/train_in_{:02d}_{:06d}.png'.format(self.samples_dir, model_dir, epoch, idx))

                    save_images(samples, [13, 13], './{:s}/{:s}/train_{:02d}_{:06d}.png'.format(self.samples_dir, model_dir, epoch, idx))
                    #save_images(batch_images, [8, 8],
                    #            './samples/input_{:02d}_{:04d}.png'.format(epoch, idx))
                    #print("[Sample] d_loss: %.8f, g_loss: %.8f" % (d_loss, g_loss))
                '''
                if np.mod(counter, 1000) == 0:
                    self.save(config.checkpoint_dir, counter)
                #if np.mod(counter, 5000) == 0 and (config.dataset == 'CASIA_MTPIE'):
                #    #with open('./{:s}/{:s}/test_recogn.txt'.format(self.checkpoint_dir, self.model_dir), "a") as myfile:
                #    self.test_IJBA(config, isReload = False, iteration = counter)

    def discriminator(self, image,  is_reuse=False, is_training = True):

        s16 = int(self.output_size/16)
        images = tf.split(axis=0, num_or_size_splits=2, value=image)
        k4_0 = []
        for ii in range(0, 2):
            if ii == 0:
                is_reuse_l = False
            else:
                is_reuse_l = True
            k0_0 = images[ii]
            k0_1 = elu(self.d_bn0_1(conv2d(k0_0, self.df_dim*1, d_h=1, d_w =1, name='d_k01_conv', reuse = is_reuse_l), train=is_training, reuse = is_reuse_l), name='d_k01_prelu')
            k0_2 = elu(self.d_bn0_2(conv2d(k0_1, self.df_dim*2, d_h=1, d_w =1, name='d_k02_conv', reuse = is_reuse_l), train=is_training, reuse = is_reuse_l), name='d_k02_prelu')
            #k0_3 =               maxpool2d(k0_2, k=2, padding='VALID')
            k1_0 = elu(self.d_bn1_0(conv2d(k0_2, self.df_dim*2, d_h=2, d_w =2, name='d_k10_conv', reuse = is_reuse_l), train=is_training, reuse = is_reuse_l), name='d_k10_prelu')
            k1_1 = elu(self.d_bn1_1(conv2d(k1_0, self.df_dim*2, d_h=1, d_w =1, name='d_k11_conv', reuse = is_reuse_l), train=is_training, reuse = is_reuse_l), name='d_k11_prelu')
            k1_2 = elu(self.d_bn1_2(conv2d(k1_1, self.df_dim*4, d_h=1, d_w =1, name='d_k12_conv', reuse = is_reuse_l), train=is_training, reuse = is_reuse_l), name='d_k12_prelu')
            #k1_3 =               maxpool2d(k1_2, k=2, padding='VALID')
            k2_0 = elu(self.d_bn2_0(conv2d(k1_2, self.df_dim*4, d_h=2, d_w =2, name='d_k20_conv', reuse = is_reuse_l), train=is_training, reuse = is_reuse_l), name='d_k20_prelu')
            k2_1 = elu(self.d_bn2_1(conv2d(k2_0, self.df_dim*3, d_h=1, d_w =1, name='d_k21_conv', reuse = is_reuse_l), train=is_training, reuse = is_reuse_l), name='d_k21_prelu')
            k2_2 = elu(self.d_bn2_2(conv2d(k2_1, self.df_dim*6, d_h=1, d_w =1, name='d_k22_conv', reuse = is_reuse_l), train=is_training, reuse = is_reuse_l), name='d_k22_prelu')
            #k2_3 =               maxpool2d(k2_2, k=2, padding='VALID')
            k3_0 = elu(self.d_bn3_0(conv2d(k2_2, self.df_dim*6, d_h=2, d_w =2, name='d_k30_conv', reuse = is_reuse_l), train=is_training, reuse = is_reuse_l), name='d_k30_prelu')
            k3_1 = elu(self.d_bn3_1(conv2d(k3_0, self.df_dim*4, d_h=1, d_w =1, name='d_k31_conv', reuse = is_reuse_l), train=is_training, reuse = is_reuse_l), name='d_k31_prelu')
            k3_2 = elu(self.d_bn3_2(conv2d(k3_1, self.df_dim*8, d_h=1, d_w =1, name='d_k32_conv', reuse = is_reuse_l), train=is_training, reuse = is_reuse_l), name='d_k32_prelu')
            #k3_3 =               maxpool2d(k3_2, k=2, padding='VALID')
            k4_0_ii = elu(self.d_bn4_0(conv2d(k3_2, self.df_dim*8, d_h=2, d_w =2, name='d_k40_conv', reuse = is_reuse_l), train=is_training, reuse = is_reuse_l), name='d_k40_prelu')
            k4_0.append(k4_0_ii)
        k4_0 = tf.concat(k4_0, 0)

        # Frontal
        k4_0s = tf.split(axis=0, num_or_size_splits=2, value=k4_0)
        k4_2 = []
        for ii in range(0, 2):
            if ii == 0:
                is_reuse_l = False
            else:
                is_reuse_l = True
            k4_1 = elu(self.d_bn4_1(conv2d(k4_0s[ii], self.df_dim*5, d_h=1, d_w =1, name='d_k41_conv', reuse = is_reuse_l), train=is_training, reuse = is_reuse_l), name='d_k41_prelu')
            k4_2_ii =     self.d_bn4_2(conv2d(k4_1, self.gfc_dim,  d_h=1, d_w =1, name='d_k42_conv', reuse = is_reuse_l), train=is_training, reuse = is_reuse_l)
            k4_2.append(k4_2_ii)
        k4_2 = tf.concat(k4_2, 0)

        k5_id = tf.nn.avg_pool(k4_2, ksize = [1, s16, s16, 1], strides = [1,1,1,1],padding = 'VALID')
        k5_id = tf.reshape(k5_id, [-1, self.dfc_dim]) #tf.nn.dropout(tf.reshape(k5_id, [-1, self.dfc_dim]), keep_prob = 0.6)

        k5_normalized = tf.nn.l2_normalize(k5_id, 0, 1e-10, name='k5_norm')
        k6_id_VGG2, w_VGG2 = linear_no_bias(k5_normalized, SUBJECT_NUM_VGG2, 'd_k6_id_VGG2', reuse=is_reuse, with_w=True)

        k6_id = k6_id_VGG2
        k6_w = w_VGG2

        # Left profile
        #k4_1 = elu(self.d_bn4_1_l(conv2d(k4_0, self.df_dim*5, d_h=1, d_w =1, name='d_k41_conv_l', reuse = is_reuse), train=is_training, reuse = is_reuse), name='d_k41_prelu_l')
        #k4_2 =     self.d_bn4_2_l(conv2d(k4_1, self.gfc_dim,  d_h=1, d_w =1, name='d_k42_conv_l', reuse = is_reuse), train=is_training, reuse = is_reuse)

        #k5 = tf.nn.avg_pool(k4_2, ksize = [1, s16, s16, 1], strides = [1,1,1,1],padding = 'VALID')
        #k5 = tf.nn.dropout(tf.reshape(k5, [-1, self.dfc_dim]), keep_prob = 0.6)
        #k6_id_l   = linear(k5, self.subject_num, 'd_k6_id_lin_l')

        # Right profile
        #k4_1 = elu(self.d_bn4_1_r(conv2d(k4_0, self.df_dim*5, d_h=1, d_w =1, name='d_k41_conv_r', reuse = is_reuse), train=is_training, reuse = is_reuse), name='d_k41_prelu_r')
        #k4_2 =     self.d_bn4_2_r(conv2d(k4_1, self.gfc_dim,  d_h=1, d_w =1, name='d_k42_conv_r', reuse = is_reuse), train=is_training, reuse = is_reuse)

        #k5 = tf.nn.avg_pool(k4_2, ksize = [1, s16, s16, 1], strides = [1,1,1,1],padding = 'VALID')
        #k5 = tf.nn.dropout(tf.reshape(k5, [-1, self.dfc_dim]), keep_prob = 0.6)
        #k6_id_r   = linear(k5, self.subject_num, 'd_k6_id_lin_r')

        # Pose
        k4_1 = elu(self.d_bn4_1_p(conv2d(k4_0, self.df_dim*5, d_h=1, d_w =1, name='d_k41_conv_pose', reuse = is_reuse), train=is_training, reuse = is_reuse), name='d_k41_prelu_pose')
        k4_2 =     self.d_bn4_2_p(conv2d(k4_1, self.dfc_dim,  d_h=1, d_w =1, name='d_k42_conv_pose', reuse = is_reuse), train=is_training, reuse = is_reuse)

        k5 = tf.nn.avg_pool(k4_2, ksize = [1, s16, s16, 1], strides = [1,1,1,1],padding = 'VALID')
        k5 = tf.nn.dropout(tf.reshape(k5, [-1, self.dfc_dim]), keep_prob = 0.6)
        k6_pose = linear(k5, self.pose_dim,    'd_k6_pose_lin', reuse = is_reuse)

        # Adversarial
        k4_1 = elu(self.d_bn4_1_a(conv2d(k4_0, self.df_dim*5, d_h=1, d_w =1, name='d_k41_conv_ad', reuse = is_reuse), train=is_training, reuse = is_reuse), name='d_k41_prelu_ad')
        k4_2 =     self.d_bn4_2_a(conv2d(k4_1, self.dfc_dim,  d_h=1, d_w =1, name='d_k42_conv_ad', reuse = is_reuse), train=is_training, reuse = is_reuse)

        k5 = tf.nn.avg_pool(k4_2, ksize = [1, s16, s16, 1], strides = [1,1,1,1],padding = 'VALID')
        k5 = tf.nn.dropout(tf.reshape(k5, [-1, self.dfc_dim]), keep_prob = 0.6)
        k6_real = linear(k5, 1, 'd_k6_ad_lin', reuse = is_reuse)

        return k6_real, k6_pose, k6_id, k5_id, k6_w#, k6_id_l, k6_id_r,

    def generator_encoder(self, image,  is_reuse=False, is_training = True):

        s16 = int(self.output_size/16)
        k0_0 = image
        k0_1 = elu(self.g_bn0_1(conv2d(k0_0, self.gf_dim*1, d_h=1, d_w =1, name='g_k01_conv', reuse = is_reuse), train=is_training, reuse = is_reuse), name='g_k01_prelu')
        k0_2 = elu(self.g_bn0_2(conv2d(k0_1, self.gf_dim*2, d_h=1, d_w =1, name='g_k02_conv', reuse = is_reuse), train=is_training, reuse = is_reuse), name='g_k02_prelu')
        #k0_3 =               maxpool2d(k0_2, k=2, padding='VALID')
        k1_0 = elu(self.g_bn1_0(conv2d(k0_2, self.gf_dim*2, d_h=2, d_w =2, name='g_k10_conv', reuse = is_reuse), train=is_training, reuse = is_reuse), name='g_k10_prelu')
        k1_1 = elu(self.g_bn1_1(conv2d(k1_0, self.gf_dim*2, d_h=1, d_w =1, name='g_k11_conv', reuse = is_reuse), train=is_training, reuse = is_reuse), name='g_k11_prelu')
        k1_2 = elu(self.g_bn1_2(conv2d(k1_1, self.gf_dim*4, d_h=1, d_w =1, name='g_k12_conv', reuse = is_reuse), train=is_training, reuse = is_reuse), name='g_k12_prelu')
        #k1_3 =               maxpool2d(k1_2, k=2, padding='VALID')
        k2_0 = elu(self.g_bn2_0(conv2d(k1_2, self.gf_dim*4, d_h=2, d_w =2, name='g_k20_conv', reuse = is_reuse), train=is_training, reuse = is_reuse), name='g_k20_prelu')
        k2_1 = elu(self.g_bn2_1(conv2d(k2_0, self.gf_dim*3, d_h=1, d_w =1, name='g_k21_conv', reuse = is_reuse), train=is_training, reuse = is_reuse), name='g_k21_prelu')
        k2_2 = elu(self.g_bn2_2(conv2d(k2_1, self.gf_dim*6, d_h=1, d_w =1, name='g_k22_conv', reuse = is_reuse), train=is_training, reuse = is_reuse), name='g_k22_prelu')
        #k2_3 =               maxpool2d(k2_2, k=2, padding='VALID')
        k3_0 = elu(self.g_bn3_0(conv2d(k2_2, self.gf_dim*6, d_h=2, d_w =2, name='g_k30_conv', reuse = is_reuse), train=is_training, reuse = is_reuse), name='g_k30_prelu')
        k3_1 = elu(self.g_bn3_1(conv2d(k3_0, self.gf_dim*4, d_h=1, d_w =1, name='g_k31_conv', reuse = is_reuse), train=is_training, reuse = is_reuse), name='g_k31_prelu')
        k3_2 = elu(self.g_bn3_2(conv2d(k3_1, self.gf_dim*8, d_h=1, d_w =1, name='g_k32_conv', reuse = is_reuse), train=is_training, reuse = is_reuse), name='g_k32_prelu')
        #k3_3 =               maxpool2d(k3_2, k=2, padding='VALID')
        k4_0 = elu(self.g_bn4_0(conv2d(k3_2, self.gf_dim*8, d_h=2, d_w =2, name='g_k40_conv', reuse = is_reuse), train=is_training, reuse = is_reuse), name='g_k40_prelu')
        k4_1 = elu(self.g_bn4_1(conv2d(k4_0, self.gf_dim*5, d_h=1, d_w =1, name='g_k41_conv', reuse = is_reuse), train=is_training, reuse = is_reuse), name='g_k41_prelu')
        k4_2 =     self.g_bn4_2(conv2d(k4_1, self.gfc_dim,  d_h=1, d_w =1, name='g_k42_conv', reuse = is_reuse), train=is_training, reuse = is_reuse)

        k5 = tf.nn.avg_pool(k4_2, ksize = [1, s16, s16, 1], strides = [1,1,1,1],padding = 'VALID')
        k5 = tf.reshape(k5, [-1, self.gfc_dim])
        if (is_training):
            k5 = tf.nn.dropout(k5, keep_prob = 0.6)

        k4_c =     self.g_bn4_c(conv2d(k4_1, 1,             d_h=1, d_w =1, name='g_h4c_conv', reuse = is_reuse), train=is_training, reuse = is_reuse)
        k5_c = tf.nn.avg_pool(k4_c, ksize = [1, s16, s16, 1], strides = [1,1,1,1],padding = 'VALID')
        k5_c = tf.nn.sigmoid(tf.reshape(k5_c, [-1, 1]))

        #k6 = linear(k5, self.subject_num, 'g_h6_real_lin')

        return k5, k5_c#, k6

    def generator_decoder(self, z, k5, pose, is_reuse=False, is_training=True):

        n_size = k5.get_shape()

        n_size = n_size[0]
        if not n_size == None:
            n_size = int(n_size)

        s = self.output_size
        s2, s4, s8, s16, s32 = int(s/2), int(s/4), int(s/8), int(s/16), int(s/32)

        # project `z` and reshape
        yb = tf.reshape(k5, [-1, self.gfc_dim])
        #pb = tf.reshape(tf.tile(pose, [1,self.pose_repeat]) , [-1, self.pose_dim*self.pose_repeat])
        pb = elu(linear(pose, self.pose_repeat, scope= 'g_h6_pose_lin', reuse = is_reuse))
        #ib = tf.reshape(pose, [-1, self.il_dim])
        zb = tf.reshape(z, [-1, self.z_dim])
        zbb = tf.concat(axis=1, values=[zb, pb, yb]) #ib
        self.z_ = linear(zbb, self.gfc_dim*s32*s32, scope= 'g_h5_lin', reuse = is_reuse)
        h5 = tf.reshape(tf.nn.dropout(self.z_, keep_prob = 0.6 ), [-1, s32, s32, self.gfc_dim])
        h5 = elu(self.g1_bn5(h5, train=is_training, reuse = is_reuse), name="g_h5_prelu")

        h4_1 = deconv2d(h5, [n_size, s16, s16, self.gf_dim*5], name='g_h4', reuse = is_reuse, with_w=False)
        h4_1 = elu(self.g1_bn4(h4_1, train=is_training, reuse = is_reuse), name="g_h4_prelu")
        h4_0 = deconv2d(h4_1, [n_size, s16, s16, self.gf_dim*8], d_h=1, d_w=1, name='g_40', reuse = is_reuse, with_w=False)
        h4_0 = elu(self.g1_bn4_0(h4_0, train=is_training, reuse = is_reuse), name="g_h40_prelu")

        h3_2 = deconv2d(h4_0, [n_size, s8, s8, self.gf_dim*8], d_h=2, d_w=2, name='g_h32', reuse = is_reuse, with_w=False)
        h3_2 = elu(self.g1_bn3_2(h3_2, train=is_training, reuse = is_reuse), name="g_h32_prelu")
        h3_1 = deconv2d(h3_2, [n_size, s8, s8, self.gf_dim*4], d_h=1, d_w=1, name='g_h31', reuse = is_reuse, with_w=False)
        h3_1 = elu(self.g1_bn3_1(h3_1, train=is_training, reuse = is_reuse), name="g_h31_prelu")
        h3_0 = deconv2d(h3_1, [n_size, s8, s8, self.gf_dim*6], d_h=1, d_w=1, name='g_h30', reuse = is_reuse, with_w=False)
        h3_0 = elu(self.g1_bn3_0(h3_0, train=is_training, reuse = is_reuse), name="g_h30_prelu")

        h2_2 = deconv2d(h3_0, [n_size, s4, s4, self.gf_dim*6], d_h=2, d_w=2, name='g_h22', reuse = is_reuse, with_w=False)
        h2_2 = elu(self.g1_bn2_2(h2_2, train=is_training, reuse = is_reuse), name="g_h22_prelu")
        h2_1 = deconv2d(h2_2, [n_size, s4, s4, self.gf_dim*3], d_h=1, d_w=1, name='g_h21', reuse = is_reuse, with_w=False)
        h2_1 = elu(self.g1_bn2_1(h2_1, train=is_training, reuse = is_reuse), name="g_h21_prelu")
        h2_0 = deconv2d(h2_1, [n_size, s4, s4, self.gf_dim*4], d_h=1, d_w=1, name='g_h20', reuse = is_reuse, with_w=False)
        h2_0 = elu(self.g1_bn2_0(h2_0, train=is_training, reuse = is_reuse), name="g_h20_prelu")

        h1_2 = deconv2d(h2_0, [n_size, s2, s2, self.gf_dim*4], d_h=2, d_w=2, name='g_h12', reuse = is_reuse, with_w=False)
        h1_2 = elu(self.g1_bn1_2(h1_2, train=is_training, reuse = is_reuse), name="g_h12_prelu")
        h1_1 = deconv2d(h1_2, [n_size, s2, s2, self.gf_dim*2], d_h=1, d_w=1, name='g_h11', reuse = is_reuse, with_w=False)
        h1_1 = elu(self.g1_bn1_1(h1_1, train=is_training, reuse = is_reuse), name="g_h11_prelu")
        h1_0 = deconv2d(h1_1, [n_size, s2, s2, self.gf_dim*2], d_h=1, d_w=1, name='g_h10', reuse = is_reuse, with_w=False)
        h1_0 = elu(self.g1_bn1_0(h1_0, train=is_training, reuse = is_reuse), name="g_h10_prelu")

        h0_2 = deconv2d(h1_0, [n_size, s, s, self.gf_dim*2], d_h=2, d_w=2, name='g_h02', reuse = is_reuse, with_w=False)
        h0_2 = elu(self.g1_bn0_2(h0_2, train=is_training, reuse = is_reuse), name="g_h02_prelu")
        h0_1 = deconv2d(h0_2, [n_size, s, s, self.gf_dim], d_h=1, d_w=1, name='g_h01', reuse = is_reuse, with_w=False)
        h0_1 = elu(self.g1_bn0_1(h0_1, train=is_training, reuse = is_reuse), name="g_h01_prelu")

        h0 = tf.nn.tanh(deconv2d(h0_1, [n_size, s, s, self.c_dim], d_h=1, d_w=1, name='g_h0', reuse = is_reuse, with_w=False))

        return h0

    def sampler(self, z, y=None, pose=None, il=None):

        k5, k5_c = self.generator_encoder(y, is_reuse=True, is_training = False)
        h0 = self.generator_decoder(z, k5, pose, is_reuse=True, is_training=False)

        return h0, k5, k5_c

    '''
    def save(self, checkpoint_dir, step):
        model_name = "DCGAN.model"
        model_dir = "%s_%s_%s_%s_%s_%s_%s" % (self.dataset_name, self.batch_size, self.output_size, self.gf_dim, self.gfc_dim, self.df_dim, self.dfc_dim)
        checkpoint_dir = os.path.join(checkpoint_dir, model_dir)
        if not os.path.exists(checkpoint_dir):
            os.makedirs(checkpoint_dir)
        self.saver.save(self.sess,
                        os.path.join(checkpoint_dir, model_name),
                        global_step=step)
    def load(self, checkpoint_dir):
        print(" [*] Reading checkpoints...")
        model_dir = "%s_%s_%s_%s_%s_%s_%s" % (self.dataset_name, self.batch_size, self.output_size, self.gf_dim, self.gfc_dim, self.df_dim, self.dfc_dim)
        checkpoint_dir = os.path.join(checkpoint_dir, model_dir)
        ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
        if ckpt and ckpt.model_checkpoint_path:
            ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
            #self.saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name))
            self.d_saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name))
            self.g_en_saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name))
            return True
        else:
            return False
    '''

    @property
    def model_dir(self):
        return "%s_%s_%s_%s_%s_%s_%s" % (self.dataset_name, self.batch_size, self.output_size, self.gf_dim, self.gfc_dim, self.df_dim, self.dfc_dim)

    def save(self, checkpoint_dir, step):
        model_name = "DCGAN.model"
        checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir)

        if not os.path.exists(checkpoint_dir):
            os.makedirs(checkpoint_dir)

        self.saver.save(self.sess, os.path.join(checkpoint_dir, model_name), global_step=step)

    def load(self, checkpoint_dir):
        import re
        print(" [*] Reading checkpoints...")
        checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir)

        ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
        if ckpt and ckpt.model_checkpoint_path:
            ckpt_name = os.path.basename(ckpt.model_checkpoint_path)

            self.d_saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name))
            self.g_en_saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name))
            #self.saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name))
            counter = int(next(re.finditer("(\d+)(?!.*\d)",ckpt_name)).group(0))
            print(" [*] Success to read {}".format(ckpt_name))

            return True, counter
        else:
            print(" [*] Failed to find a checkpoint")

            return False, 0


from __future__ import print_function
import sys
import numpy as np
import geometry_constants
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from matplotlib import collections as mc

# counter geometry: top left ([0][0]) is at (0,0) in the dimension
# viewed facing the indicated end
# "trans" configure:
#  __________________
# |_____|_____|_____|___...
#    |_____|_____|_____|___...
#       |_____|_____|_____|___...
#          |_____|_____|_____|   ...
#
# "cis" configure:
#              __________________
#          ___|_____|_____|_____|...
#      ___|_____|_____|_____|...
#  ___|_____|_____|_____|...
# |_____|_____|_____|...
# FIXME: Further 'dimensions' can be added for vertically placed modules

def dimensionGen(type, ndc = geometry_constants.dicounterPerLayer, nl = geometry_constants.nLayer):
    if type != 'trans' and type != 'cis':
        sys.exit("Error: dimensionGen: 'type' cannot be " + type)
    offsetSign = 1 if type == 'trans' else -1
    _chanY = np.array([[-geometry_constants.scintillatorBarThickness*j-sum(geometry_constants.gapBetweenLayers[:j])for i in range(4*ndc)] for j in range(nl)])
    _chanX = np.array([[geometry_constants.xFiberPosition[i%4]+geometry_constants.dicounterTranslationalX*int(i/4)+geometry_constants.layerOffset*j*offsetSign for i in range(4*ndc)] for j in range(nl)])
    dimension = {"channelX": _chanX,
                 "channelY": _chanY,
                 "cntrX1":_chanX[:, ::2]-geometry_constants.channelToEdge,
                 "cntrX2":_chanX[:, 1::2]+geometry_constants.channelToEdge,
                 "cntrY1":_chanY[:, ::2]-geometry_constants.scintillatorBarThickness/2.,
                 "cntrY2":_chanY[:, 1::2]+geometry_constants.scintillatorBarThickness/2.}
    return dimension

class testBenchGeometry:
    # object containing coordinates of channels, extrusions, mapping

    def __init__(self, geom_dict):

        # self.tag: a string used as a label
        # self.FEBList = [] list of relevant FEBs
        # self.coordinates = {"channelX":..., "channelY":..., "cntrX1":..., "cntrX2":..., "cntrY1":..., "cntrY2":...}
        # self.mappingFEBCh = ...
        # self.idleMask = ... ; boolean type indicating no FEB connection
        # self.badChMask = ... ; boolean type, bad channels with FEB connection
        # self.analysisMask = ... ; boolean type, channels with FEB connection, designated for analysis
        # self.triggerOnlyMask = ... ; boolean type, channels with FEB connection, designated for triggering only
        # self.canvasSize = (x,y), 2-tuple in inches
        # ...: 3-D array of dimension nModule*nLayer*nFiber(or mCounter)

        if not isinstance(geom_dict, dict):
            sys.exit("Error: testBenchGeometry: initialization requires a dict type argument")
        if not {'module','moduleOffsetX','moduleOffsetY','FEB'}.issubset(set(geom_dict.keys())):
            sys.exit("Error: testBenchGeometry: initialization requires a dict containing {'module','moduleOffsetX','moduleOffsetY','FEB'}")
        tLen = len(geom_dict['module'])
        for item in ['module','moduleOffsetX','moduleOffsetY','FEB']:
            if len(geom_dict[item])!=tLen:
                sys.exit("Error: testBenchGeometry: initialization requires a dict with values of same length")
        self.tag = geom_dict['tag']
        self.FEBList = []
        for iList in geom_dict['FEB']:
            for iFEB in iList:
                if iFEB not in self.FEBList and iFEB!=-1:
                    self.FEBList.append(iFEB)

        dimensionNameList = ["channelX", "channelY", "cntrX1", "cntrX2", "cntrY1", "cntrY2"]
        self.coordinates = {}
        for item in dimensionNameList:
            self.coordinates.update({item: []})
        self.mappingFEBCh = []
        for index, module in enumerate(geom_dict['module']):
            if module not in geometry_constants.module_geom_dict.keys():
                sys.exit("Error: testBenchGeometry: %s is not a defined module geometry"%(module))
            tDimension = dimensionGen(geometry_constants.module_geom_dict[module]['type'])
            for item in dimensionNameList:
                tOffset = geom_dict['moduleOffsetX'][index] if 'X' in item else geom_dict['moduleOffsetY'][index]
                self.coordinates[item].append(tOffset + tDimension[item])
            tMap = geometry_constants.module_geom_dict[module]['map'].copy()
            with np.nditer(tMap, op_flags=['readwrite']) as handle:
                for x in handle:
                    if int(x.real) not in range(len(geom_dict['FEB'][index])):
                        sys.exit("Error: testBenchGeometry: 'FEB'[%i] needs to have a length of at least %i; fill in the blank with -1"%(index, 1+int(x.real)))
                    x[...] = 1j*x.imag + geom_dict['FEB'][index][int(x.real)]
            self.mappingFEBCh.append(tMap)

        for item in dimensionNameList:
            self.coordinates[item] = np.array(self.coordinates[item])
        self.mappingFEBCh = np.array(self.mappingFEBCh)

        self.idleMask = (self.mappingFEBCh.real == -1.)

        self.badChMask = (self.mappingFEBCh == None)
        if 'badChannels' in geom_dict.keys():
            if geom_dict['badChannels']:
                for tTuple in geom_dict['badChannels']:
                    self.badChMask[tuple(np.argwhere(self.mappingFEBCh==tTuple[0]+1j*tTuple[1])[0])] = True

        self.analysisMask = (~self.idleMask)
        if 'triggerOnlyChs' in geom_dict.keys():
            if geom_dict['triggerOnlyChs']:
                for tTuple in geom_dict['triggerOnlyChs']:
                    self.analysisMask[tuple(np.argwhere(self.mappingFEBCh==tTuple[0]+1j*tTuple[1])[0])] = False
        self.triggerOnlyMask = ((~self.idleMask)&(~self.analysisMask))

        self.canvasSize = (8,6)
        if 'canvasSize' in geom_dict.keys():
            self.canvasSize = geom_dict['canvasSize']

    def chToFiberIndex(self, iFEB, iCh):
        # converts (iFEB, iCh) to a tuple that can access the coordinates of the corresponding fiber
        return tuple(np.argwhere(self.mappingFEBCh==iFEB+1j*iCh)[0])

    def chToCntrIndex(self, iFEB, iCh):
        # converts (iFEB, iCh) to a tuple that can access the coordinates of the corresponding counter
        tTuple = tuple(np.argwhere(self.mappingFEBCh==iFEB+1j*iCh)[0])
        tTuple = tuple(tTuple[0], tTuple[1], int(tTuple[2]/2))
        return tTuple

    def plotCntrCoordGen(self, mask):
        tchtrX1_masked = self.coordinates['cntrX1'][mask] # this returns flattened array
        tchtrX2_masked = self.coordinates['cntrX2'][mask]
        tchtrY1_masked = self.coordinates['cntrY1'][mask]
        tchtrY2_masked = self.coordinates['cntrY2'][mask]
        #tchtrX1_masked = np.ma.masked_where(~mask, self.coordinates['cntrX1']) # converting masked elements to NaN causes warnings
        #tchtrX2_masked = np.ma.masked_where(~mask, self.coordinates['cntrX2'])
        #tchtrY1_masked = np.ma.masked_where(~mask, self.coordinates['cntrY1'])
        #tchtrY2_masked = np.ma.masked_where(~mask, self.coordinates['cntrY2'])

        lineSegList = []
        if tchtrX1_masked.any():
            with np.nditer(tchtrX1_masked, flags=['multi_index']) as it:
                for x in it:
                    lineSegList.append([(tchtrX1_masked[it.multi_index], tchtrY1_masked[it.multi_index]),
                                        (tchtrX1_masked[it.multi_index], tchtrY2_masked[it.multi_index]),
                                        (tchtrX2_masked[it.multi_index], tchtrY2_masked[it.multi_index]),
                                        (tchtrX2_masked[it.multi_index], tchtrY1_masked[it.multi_index]),
                                        (tchtrX1_masked[it.multi_index], tchtrY1_masked[it.multi_index])])
        return lineSegList

    def plotGeometry(self, ax, titleExtra = '', dotChannel = False, verbose = True):
        idleCounterMask = (self.idleMask[:, :, 1::2])&(self.idleMask[:, :, ::2])
        linesCoord_solid = self.plotCntrCoordGen(~idleCounterMask)
        linesCoord_dotted = self.plotCntrCoordGen(idleCounterMask)
        if linesCoord_solid:
            lc_solid = mc.LineCollection(linesCoord_solid, linewidths=0.5, linestyles='solid',
                                         colors='#080808',alpha=0.8)
            ax.add_collection(lc_solid)
        if linesCoord_dotted:
            lc_dotted = mc.LineCollection(linesCoord_dotted, linewidths=0.5, linestyles='dotted',
                                         colors='#434343',alpha=0.8)
            ax.add_collection(lc_dotted)

        if dotChannel:
            activeChMask = (~self.idleMask)&(~self.badChMask)
            plt.scatter(self.coordinates['channelX'][activeChMask],
                        self.coordinates['channelY'][activeChMask],
                        c='#080808', alpha=0.8, marker='.')

        crossMask = self.badChMask|(self.idleMask!=np.repeat(idleCounterMask, 2, axis=2))
        plt.scatter(self.coordinates['channelX'][crossMask],
                    self.coordinates['channelY'][crossMask],
                    c='#434343', alpha=0.8, marker='x')

        ax.set_title(self.tag+' '+titleExtra)
        ax.set_xlabel('x [mm]')
        ax.set_ylabel('y [mm]')

        if verbose:
            with np.nditer(self.mappingFEBCh, flags=['multi_index']) as it:
                for x in it:
                    if (self.mappingFEBCh[it.multi_index].imag in [0,31,32, 63]) and self.mappingFEBCh[it.multi_index].real>=0:
                        ax.text(self.coordinates['channelX'][it.multi_index],self.coordinates['channelY'][it.multi_index]-2,
                                '%i_%i'%(self.mappingFEBCh[it.multi_index].real,self.mappingFEBCh[it.multi_index].imag),
                                color='black',verticalalignment='top', horizontalalignment='center',zorder = 5,fontsize = 3)

        ax.autoscale()
        plt.axis('equal')
        return


from leather.series.base import Series
from leather.utils import Datum, X, Y, Z

class CategorySeries(Series):
    """
    A series of categorized data and its associated metadata.

    Series object does not modify the data it is passed.

    :param data:
        A sequence (rows) of sequences (columns), a.k.a. :func:`csv.reader`
        format. If the :code:`x` and :code:`y` are not specified then the first
        column is used as the X values and the second column is used for Y.

        Or, a sequence of (rows) of dicts (columns), a.k.a.
        :class:`csv.DictReader` format. If this format is used then :code:`x`
        and :code:`y` arguments must specify the columns to be charted.

        Or, a custom data format, in which case :code:`x` and :code:`y` must
        specify :func:`.key_function`.
    :param x:
        If using sequence row data, then this may be either an integer index
        identifying the X column, or a :func:`.key_function`.

        If using dict row data, then this may be either a key name identifying
        the X column, or a :func:`.key_function`.

        If using a custom data format, then this must be a
        :func:`.key_function`.`
    :param y:
        See :code:`x`.
    :param z:
        See :code:`y`. This variable identifies the category/sub-series of each
        row.
    :param name:
        An optional name to be used in labeling this series. This will be
        used as the chart title if rendered in a :class:`.Lattice`.
    """
    def __init__(self, data, x=None, y=None, z=None, name=None):
        self._data = data
        self._name = name

        self._keys = [
            self._make_key(x if x is not None else X),
            self._make_key(y if y is not None else Y),
            self._make_key(z if z is not None else Z)
        ]

        self._types = [
            self._infer_type(X),
            self._infer_type(Y),
            self._infer_type(Z)
        ]

    def data(self):
        """
        Return data for this series grouped for rendering.
        """
        x = self._keys[X]
        y = self._keys[Y]
        z = self._keys[Z]

        for i, row in enumerate(self._data):
            yield Datum(i, x(row, i), y(row, i), z(row, i), row)

    def categories(self):
        """
        Return all unique values in the category field.
        """
        z = self._keys[Z]
        categories = []

        for i, row in enumerate(self._data):
            cat = z(row, i)

            if cat not in categories:
                categories.append(cat)

        return categories


# list的拷贝
# 对列表和字典，赋值只是引用
# 对字符串，整型，元组就是保存本身

import copy

print("是引用，一变全变")
str1 = [1,2,3,4]
str2 = str1  #是引用，一变全变

str1[1]=99
print (str1,str2)

print("是复制，各自独立")
str1 = [1,2,3,4]
str2 = copy.copy(str1)  #是复制，各自独立
str1[1]=99
print (str1,str2)

print("list第一级是复制，2级及以下是引用")
str1 = [1,2,3,[4,5,[6,7,8]]]
str2 = copy.copy(str1)  #list的第一级是复制，各自独立，2级及以下是引用
str1[1]=99
str1[3][1]=99
print (str1,str2)

print("list各级都是复制，各自独立")
str1 = [1,2,3,[4,5,[6,7,8]]]
str2 = copy.deepcopy(str1)  #是复制，各自独立
str1[1]=99
str1[3][1]=99
print (str1,str2)


import cv2
import numpy as np
import os
import matplotlib.pyplot as plt
import time
import math
import matplotlib.animation as animation
import os
import glob
from natsort import natsorted

def remove_noise(image, kernel_size):
    return cv2.GaussianBlur(image, (kernel_size),0)

def filter_region(image, vertices):
    mask = np.zeros_like(image)
    if len(mask.shape)==2:
        cv2.fillPoly(mask, vertices, 255)
    else:
        cv2.fillPoly(mask, vertices, (255,)*mask.shape[2])
    return cv2.bitwise_and(image, mask)

DATA_DIR = "C:\\Users\\luluk\\OneDrive - imt-atlantique\\Imt_courses\\PROCOM\\ia_racing_imt-main\\supervise\\dataset_drive8\\images"
IM_DIR = "C:\\Users\\luluk\\OneDrive - imt-atlantique\\Imt_courses\\PROCOM\\ia_racing_imt-main\\supervise\\dataset_drive4\\images_preprocessed"
# Create the new image folder if it doesn't exist
if not os.path.exists(IM_DIR):
    os.makedirs(IM_DIR)

# Load each image and store them in an array
images = []
images = [cv2.cvtColor(cv2.imread(file),cv2.COLOR_BGR2RGB) for file in natsorted(glob.glob(DATA_DIR+"\\*.jpg"))]
# for file in image_files:
#     image = cv2.imread(file, cv2.IMREAD_COLOR)
#     images.append(image)

#######################################################
# Now you have an array of images from the data directory
plt.imshow(images[0], animated=True)
plt.show(block=False)
plt.pause(0.005)
plt.clf()
#######################################################
# Define the source and destination points for the perspective transform
xs=[10, 150, 118, 35]
ys=[69,73,60,60]
src = np.float32([(xs[0], ys[0]), (xs[1], ys[1]), (xs[2], ys[2]), (xs[3], ys[3])])
# Destination points are to be parallel, taking into account the image size
dst = np.float32([[40, 120],
                [120, 120],
                [120, 0],
                [40, 0]])

# Calculate the transformation matrix and it's inverse transformation
M = cv2.getPerspectiveTransform(src, dst)
M_inv = cv2.getPerspectiveTransform(dst, src)

#######################################################
TRESH_MIN = np.array([0, 75, 0],np.uint8)
TRESH_MAX = np.array([180, 255, 20],np.uint8)
TRESHY_MIN = np.array([0, 210, 210],np.uint8)
TRESHY_MAX = np.array([200, 255, 255],np.uint8)
TRESH3_MIN = np.array([0, 0, 0],np.uint8)
TRESH3_MAX = np.array([120, 180, 180],np.uint8)
#######################################################
verticesx=np.array([[(0, 60), (160, 60), (160, 120), (0, 120)]])

################################
# Lane detection configuration #
################################
nwindows = 40
window_height = math.floor(images[0].shape[0]//nwindows)

# Set the width of the windows +/- margin
margin = 5

# Set minimum number of pixels found to recenter window
minpix = 6

#######################################################
# Define conversions in x and y from pixels space to meters
ym_per_pix = 0.18/120 # meters per pixel in y dimension
xm_per_pix = 0.5/80 # meters per pixel in x dimension

#######################################################
video = cv2.VideoWriter("video_CV.avi", 0, 20, (480,360))

#######################################################
leftx_current = 40
rightx_current = 120
#######################################################
radstock = []

start_time = time.time()

p=0
r=0
l=0
#######################################################
for i in range(len(images)):
    imagex=images[i]
    ##########################################################
    # Convert to grayscale, blur, and perform edge detection #
    ##########################################################
    test_image = imagex
    #test_image = remove_noise(image=test_image, kernel_size=(1, 1))
    test_image = cv2.warpPerspective(test_image, M, (160,120))
    # test_image = cv2.cvtColor(test_image, cv2.COLOR_BGR2HLS)
    test_image = 255-cv2.inRange(test_image, TRESH3_MIN, TRESH3_MAX)
    # #stock = test_image
    # #test_image = test_image-cv2.inRange(imagex, TRESHY_MIN, TRESHY_MAX)
    # #test_image = filter_region(image=test_image,vertices=verticesx)
    # # #stock =filter_region(image=stock,vertices=verticesx)

    # # #########################
    # # # Perspective transform #
    # # #########################
    # # test_image = cv2.warpPerspective(test_image, M, (160,120))

    # #test_image = cv2.rotate(test_image, cv2.ROTATE_180)
    # # #stock = cv2.warpPerspective(stock, M, (160,120))

    ####################################################
    # Take a histogram of the bottom half of the image #
    ####################################################
    try:

        histogram = np.sum(test_image[0:10,:], axis=0)
        histogrambis = np.sum(test_image[100:120,:], axis=0)
        midpoint = math.floor((histogram.nonzero()[0][0]+histogram.nonzero()[0][-1])//2)
        ###############################################################
        # Find the peak of the left and right halves of the histogram #
        ###############################################################
        if i==0 or histogram[max(leftx_base-margin,0):min(leftx_base+margin,160)+1]==0 or histogram[max(rightx_base-margin,0):min(rightx_base+margin,160)+1]==0 or histogram.nonzero()[0][-1]-histogram.nonzero()[0][0]>85 or l>=5 or r>=5:
            leftx_base = midpoint-histogram[:midpoint][::-1].nonzero()[0][0]-3
            rightx_base = np.argmax(histogram[midpoint:]) + midpoint + 3
            if rightx_base-leftx_base<60:
                if np.mean(histogrambis.nonzero()[0])>80:
                    leftx_base=0
                else:
                    rightx_base=160
    except:
        pass

    # Set height of windows
    leftx_current = leftx_base
    rightx_current = rightx_base
    nonzeroy = np.array(test_image.nonzero()[0])
    nonzerox = np.array(test_image.nonzero()[1])
    if i==0 or ((left_lanex_inds.shape==(0,) and left_laney_inds.shape == (0,)) and (right_lanex_inds.shape == (0,) and right_laney_inds.shape == (0,))):
        left_lanex_inds = []
        right_lanex_inds = []
        left_laney_inds = []
        right_laney_inds = []
        try:
            for window in range(nwindows):
                # Identify window boundaries in x and y (and right and left)
                win_y_high = (window+1)*window_height
                win_y_low = window*window_height

                win_xleft_low = math.floor(max(leftx_current - margin,0))
                win_xleft_high = math.floor(min(leftx_current + margin,test_image.shape[1]))
                win_xright_low = math.floor(max(rightx_current - margin,0))
                win_xright_high = math.floor(min(rightx_current + margin,test_image.shape[1]))

                if win_xleft_low>=(win_y_low-25)/3.1:
                    # Identify the nonzero pixels in x and y within the window #
                    good_lefty_inds = nonzeroy[((nonzeroy >= win_y_low) & (nonzeroy < win_y_high) &
                    (nonzerox >= win_xleft_low) &  (nonzerox < win_xleft_high))]
                    good_leftx_inds = nonzerox[((nonzeroy >= win_y_low) & (nonzeroy < win_y_high) &
                    (nonzerox >= win_xleft_low) &  (nonzerox < win_xleft_high))]

                    # Append these indices to the lists
                    left_laney_inds.append(good_lefty_inds)
                    left_lanex_inds.append(good_leftx_inds)

                    # If you found > minpix pixels, recenter next window on their mean position
                    if len(good_leftx_inds) > minpix:
                        leftx_current = math.ceil(np.mean(good_leftx_inds))

                if win_xright_high<=(700-win_y_low)/4.7:
                    good_righty_inds = nonzeroy[((nonzeroy >= win_y_low) & (nonzeroy < win_y_high) &
                    (nonzerox >= win_xright_low) &  (nonzerox < win_xright_high))]
                    good_rightx_inds = nonzerox[((nonzeroy >= win_y_low) & (nonzeroy < win_y_high) &
                    (nonzerox >= win_xright_low) &  (nonzerox < win_xright_high))]

                    #Append these indices to the lists
                    right_laney_inds.append(good_righty_inds)
                    right_lanex_inds.append(good_rightx_inds)

                    # If you found > minpix pixels, recenter next window on their mean position
                    if len(good_rightx_inds) > minpix:
                        rightx_current = math.ceil(np.mean(good_rightx_inds))

            # Concatenate the arrays of indices
            left_laney_inds = np.concatenate(left_laney_inds)
            right_laney_inds = np.concatenate(right_laney_inds)
            left_lanex_inds = np.concatenate(left_lanex_inds)
            right_lanex_inds = np.concatenate(right_lanex_inds)
        except:
            left_laney_inds = []
            right_laney_inds = []
            left_lanex_inds = []
            right_lanex_inds = []
    else:
        try:
            left_lane_inds = ((nonzerox > (left_fit[0]*(nonzeroy**2) + left_fit[1]*nonzeroy +
                        left_fit[2] - margin)) & (nonzerox < (left_fit[0]*(nonzeroy**2) +
                        left_fit[1]*nonzeroy + left_fit[2] + margin)))
            left_lanex_inds = nonzerox[left_lane_inds]
            left_laney_inds = nonzeroy[left_lane_inds]
        except:
            left_lane_inds = []
            left_lanex_inds = []
            left_laney_inds = []
        try:
            right_lane_inds = ((nonzerox > (right_fit[0]*(nonzeroy**2) + right_fit[1]*nonzeroy +
                        right_fit[2] - margin)) & (nonzerox < (right_fit[0]*(nonzeroy**2) +
                        right_fit[1]*nonzeroy + right_fit[2] + margin)))
            right_lanex_inds = nonzerox[right_lane_inds]
            right_laney_inds = nonzeroy[right_lane_inds]

        except:
            right_lane_inds = []
            right_lanex_inds = []
            right_laney_inds = []

    try:
        left_fit = np.polyfit(left_laney_inds, left_lanex_inds, 2)
        ltrue=True
        l=0
    except:
        left_fit = [0,0,0]
        ltrue=False
        l+=1
    try:
        right_fit = np.polyfit(right_laney_inds,right_lanex_inds, 2)
        rtrue=True
        r=0
    except:
        right_fit = [0,0,0]
        rtrue=False
        r+=1

    # plt.imshow(cv2.warpPerspective(imagex,M,(160,120)))
    if rtrue:
        try:
            xrs = np.linspace(0, test_image.shape[1], 100)
            yrs = right_fit[0] * xrs ** 2 + right_fit[1] * xrs + right_fit[2]
            y_eval_r = np.max(yrs)
            y_max = test_image.shape[0]
            right_fit_cr = np.polyfit(yrs*ym_per_pix, xrs*xm_per_pix, 2)
            right_curverad = ((1 + (2*right_fit_cr[0]*y_eval_r*ym_per_pix + right_fit_cr[1])**2)**1.5) / (2*right_fit_cr[0])
            right_x_pos = right_fit[0]*y_max**2 + right_fit[1]*y_max + right_fit[2]
        except:
            right_fit_cr = [0,0,0]
            right_curverad = np.inf
            right_x_pos = np.inf
            xrs = []
            yrs = []
    else:
        right_fit_cr = [0,0,0]
        right_curverad = np.inf
        left_x_pos = np.inf
        xrs = []
        yrs = []

    if ltrue:
        try:
            xls = np.linspace(0, test_image.shape[1], 100)
            yls = left_fit[0] * xls ** 2 + left_fit[1] * xls + left_fit[2]
            y_eval_l = np.max(yls)
            y_max = test_image.shape[0]
            left_fit_cr = np.polyfit(yls*ym_per_pix, xls*xm_per_pix, 2)
            left_curverad = ((1 + (2*left_fit_cr[0]*y_eval_l*ym_per_pix + left_fit_cr[1])**2)**1.5) / (2*left_fit_cr[0])
            left_x_pos = left_fit[0]*y_max**2 + left_fit[1]*y_max + left_fit[2]
        except:
            left_fit_cr = [0,0,0]
            left_curverad = np.inf
            left_x_pos = np.inf
            xls = []
            yls = []
    else:
        left_fit_cr = [0,0,0]
        left_curverad = np.inf
        right_x_pos = np.inf
        xls = []
        yls = []

    # Calculate the x position of the center of the lane
    if left_x_pos != np.inf and right_x_pos != np.inf:
        center_lanes_x_pos = (left_x_pos + right_x_pos)//2

    elif left_x_pos != np.inf:
        center_lanes_x_pos = left_x_pos+42

    elif right_x_pos != np.inf:
        center_lanes_x_pos = right_x_pos-42

    # Calculate the deviation between the center of the lane and the center of the picture
    # The car is assumed to be placed in the center of the picture
    # If the deviation is negative, the car is on the felt hand side of the center of the lane
    veh_pos = ((test_image.shape[1]//2) - center_lanes_x_pos) * xm_per_pix

    warp_zero = np.zeros_like(test_image).astype(np.uint8)
    color_warp = np.dstack((warp_zero, warp_zero, warp_zero))
    line = np.dstack((warp_zero, warp_zero, warp_zero))
    linebis = np.dstack((warp_zero, warp_zero, warp_zero))

    if ltrue:
        pts_left_lane = np.array([(np.vstack([np.array(yls)+2.5, np.array(xls)])).T],dtype=np.int32)
        pts_left = np.array([(np.vstack([np.array(yls), np.array(xls)])).T],dtype=np.int32)
    else:
        pts_left = np.array([],dtype=np.int32)
        pts_left_lane = np.array([],dtype=np.int32)

    if rtrue:
        pts_right_lane = np.array([np.flipud((np.vstack([np.array(yrs)-2.5, np.array(xrs)])).T)],dtype=np.int32)
        pts_right = np.array([np.flipud((np.vstack([np.array(yrs), np.array(xrs)])).T)],dtype=np.int32)
    else:
        pts_right = np.array([],dtype=np.int32)
        pts_right_lane = np.array([],dtype=np.int32)

    try:
        pts_lane = np.hstack((pts_left_lane, pts_right_lane))
        cv2.fillPoly(color_warp, pts_lane.round(0), (0, 255, 0))
    except:
        pass
    if ltrue:
        cv2.polylines(line, pts_left.round(0), isClosed=False, color=(255, 0, 0), thickness=5)
        cv2.polylines(linebis, pts_left.round(0), isClosed=False, color=(255, 255, 255), thickness=5)

    if rtrue:
        cv2.polylines(line, pts_right.round(0), isClosed=False, color=(255, 0, 0), thickness=5)
        cv2.polylines(linebis, pts_right.round(0), isClosed=False, color=(255, 255, 255), thickness=5)

    tosave = color_warp | linebis
    # Save the image in the new folder
    nametosave = str(str(i)+"_cam_image_array_.jpg")
    #cv2.imwrite(os.path.join(IM_DIR, nametosave), tosave)

    # Warp the blank back to original image space using inverse perspective matrix (M_inv)
    newwarp = cv2.warpPerspective(color_warp, M_inv, (imagex.shape[1], imagex.shape[0]))
    newwarpline = cv2.warpPerspective(line, M_inv, (imagex.shape[1], imagex.shape[0]))
    newwarpmask = cv2.warpPerspective(linebis, M_inv, (imagex.shape[1], imagex.shape[0]))
    # Combine the result with the original image
    out_img = imagex
    mask = newwarpmask & imagex
    out_img = out_img - mask
    out_img = cv2.addWeighted(out_img, 1, newwarpline, 1, 0)
    out_img = cv2.addWeighted(out_img, 1, newwarp, 0.4, 0)
    out_img = cv2.line(out_img, (80, 70), (80, 120), (255, 0, 0), 1)

    resized = cv2.resize(out_img, (480,360), interpolation = cv2.INTER_LANCZOS4)
    rad = (left_curverad+right_curverad)/2
    if abs(left_curverad)>=1.2 and abs(right_curverad)<1.2:
        rad=np.sign(right_curverad)*0.25+right_curverad
    elif abs(right_curverad)>=1.2 and abs(left_curverad)<1.2:
        rad=np.sign(left_curverad)*0.25+left_curverad
    elif abs(left_curverad)>=1.2 and abs(right_curverad)>=1.2:
        rad=np.inf
    # if rad!=np.inf:
    #     if len(radstock)>=3:
    #         radstock.pop(0)
    #     radstock.append(rad)
    #     radmean=np.mean(radstock)
    #     k=0
    # elif k>=5:
    #     radmean=np.inf
    #     k+=1
    # else:
    #     radmean=np.mean(radstock)
    #     k+=1
    cv2.putText(resized,'Curve Radius [m]: '+ str(np.round(rad,2)),(10,15), cv2.FONT_HERSHEY_PLAIN, 1, (255,255,255),1,cv2.LINE_AA)
    cv2.putText(resized,'Center Offset [m]: '+ str(np.round(veh_pos,2)),(10,30), cv2.FONT_HERSHEY_PLAIN, 1,(255,255,255),1,cv2.LINE_AA)
    cv2.putText(resized,'Lfit: '+str(np.round(left_fit_cr[0],2)),(10,45), cv2.FONT_HERSHEY_PLAIN, 1,(255,255,255),1,cv2.LINE_AA)
    cv2.putText(resized,'Rfit: '+str(np.round(right_fit_cr[0],2)),(10,60), cv2.FONT_HERSHEY_PLAIN, 1,(255,255,255),1,cv2.LINE_AA)
    # cv2.putText(resized,'image n '+str(i),(10,75), cv2.FONT_HERSHEY_PLAIN, 1,(255,255,255),1,cv2.LINE_AA)
    plt.imshow(resized,animated=True)
    plt.draw()
    plt.pause(0.1-time.time()+start_time)
    plt.clf()
    plt.title('Fps:'+str(np.round(1/(time.time()-start_time),1)),loc='left')
    start_time = time.time()
    #video.write(resized)

video.release()


import requests
from bs4 import BeautifulSoup

url = "https://news.ycombinator.com/"
response = requests.get(url)

if response.status_code == 200:
    soup = BeautifulSoup(response.text, 'html.parser')
    headings = soup.select(".titleline a")

    for heading in headings:
        print(heading.text)
else:
    print("Erro ao acessar a página")


orders_count = int(input("Введите количество заказов: "))
orders = dict()

for order_number in range(1, orders_count+1):
    order = input("{}-й заказ: ".format(order_number)).split()
    if order[0] in orders:
        if order[1] in orders[order[0]]:
            orders[order[0]][order[1]] += order[2]
        else:
            orders[order[0]][order[1]] = order[2]
    else:
        orders[order[0]] = {order[1]: order[2]}

print("Итоговый заказ:")

for name in orders.keys():
    print(name, ":")
    for order in orders[name]:
        print(order, ":", orders[name][order])
    print("")


"""This class holds all the tests to test the API_Manager requests, and request processing"""

import API_Manager
import unittest

class TestAPIPRocessing(unittest.TestCase):

    def test_get_routes_response_processing(self):
        """Test that the API_Manager.get_routes response is processed into a list of route numbers correctly"""

        example_response = [
            { "Route": "4", "sdfdgsg": "Sdfsdfsf"},
            { "Route": "40", "asduhsdif": "Sdfsdfsf"},
            { "Route": "B", "sdfijsdfoij": "Sdfsdfsf"},
        ]

        route_numbers = API_Manager.route_numbers_from_response(example_response)
        self.assertEqual(['4', '40', 'B'], route_numbers)

    def test_get_directions_response_processing(self):
        """Testing that the API_Manager.get_directions processes the directions into a list correctly"""

        example_response = [{"Text": "NORTHBOUND", "Value": "4"},
                            {"Text": "SOUTHBOUND", "Value": "1"}]

        directions = API_Manager.directions_from_response(example_response)
        self.assertEqual(['NORTHBOUND', 'SOUTHBOUND'], directions)

    def test_get_Stops_response_processing(self):
        """Testing that the API_Manager.get_stops processes the stop information into a dictionary"""

        example_response = [{"Text":"Mall of America Transit Station","Value":"MAAM"},
                            {"Text":"Portland Ave and 77th St","Value":"77PO"},
                            {"Text":"Portland Ave and 66th St","Value":"66PO"}]

        stop_dict = API_Manager.stops_from_response(example_response)
        self.assertEqual({"Mall of America Transit Station": "MAAM",
                          "Portland Ave and 77th St": "77PO",
                          "Portland Ave and 66th St": "66PO"}, stop_dict)

    def test_get_times_response_processing(self):
        """Testing the API_Manager.get_times processes the times into a list"""

        example_response = [{"BlockNumber":1058,"DepartureText":"3 Min","DepartureTime":"\/Date(1557603720000-0500)\/","Route":"5","RouteDirection":"NORTHBOUND"},
                            {"BlockNumber":1053,"DepartureText":"8 Min","DepartureTime":"\/Date(1557604020000-0500)\/","Route":"5","RouteDirection":"NORTHBOUND"},
                            {"BlockNumber":1055,"DepartureText":"2:59","DepartureTime":"\/Date(1557604740000-0500)\/","Route":"5","RouteDirection":"NORTHBOUND"}
                            ]

        times = API_Manager.times_from_response(example_response)
        self.assertEqual(['3 Min', '8 Min', '2:59'], times)

    def test_build_marker_list_processing(self):
        """Testing API_Manager.build_marker_list for proper string concatenation"""

        example_list = ['Apple Valley', 'Minneapolis']

        list = API_Manager.build_marker_list(example_list)
        self.assertEqual(['size:small|color:red|Apple Valley', 'size:small|color:red|Minneapolis'], list)

    def test_build_visible_list_processing(self):
        """Testing API_Manager.build_visible_list for proper list conversion"""

        example_list = ['Apple Valley', 'Bloomington', 'St. Paul', 'Minneapolis']

        list = API_Manager.build_visible_list(example_list)
        self.assertEqual(['Apple Valley', 'Minneapolis'], list)

    def test_convert_to_direction_code(self):
        """Test that convert_to_direction_code returns the correct code"""

        direction = 'NORTHBOUND'
        code = API_Manager.convert_to_direction_code(direction)
        self.assertEqual('4', code)

        direction = 'SOUTHBOUND'
        code = API_Manager.convert_to_direction_code(direction)
        self.assertEqual('1', code)

        direction = 'EASTBOUND'
        code = API_Manager.convert_to_direction_code(direction)
        self.assertEqual('2', code)

        direction = 'WESTBOUND'
        code = API_Manager.convert_to_direction_code(direction)
        self.assertEqual('3', code)

if __name__ == '__main__':
    unittest.main()