########################################################################################################################
#                                          N-HANS speech separator                                                     #
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
#   Description:      N-HANS contains two systems, speech denoiser and speech separator. This is the first part.       #
#   Authors:          Shuo Liu, Gil Keren, Bjoern Schuller                                                             #
#   Affiliation:      Chair of Embedded Intelligence for Health Care and Wellbeing, University of Augsburg (UAU)  #
#   Version:          2.0                                                                                              #
#   Last Update:      May. 06, 2020                                                                                    #
#   Dependence Files: reader.py  blocks.py                                                                             #
#   Contact:          shuo.liu@informatik.uni-augburg.de                                                               #
########################################################################################################################
#                                                                                                                      #
#   Copyright (C) 2019, Shuo Liu, Gil Keren and Björn Schuller: University of Augsburg.                                #
#                                                                                                                      #
# This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public    #
# License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any       #
# later version.                                                                                                       #
#                                                                                                                      #
# This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied   #
# warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.#
#                                                                                                                      #
# You should have received a copy of the GNU General Public License along with this program. #  If not, see            #
# <http://www.gnu.org/licenses/>.                                                                                      #
#                                                                                                                      #
########################################################################################################################

# Import from standard libraries
from __future__ import division, absolute_import, print_function
import os, time, numpy as np, scipy.io.wavfile, functools
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import warnings
warnings.filterwarnings('ignore', category=FutureWarning)
import tensorflow as tf
if int(tf.__version__.split('.')[0]) == 2:
    import tensorflow.compat.v1 as tf
    tf.compat.v1.disable_v2_behavior()
from tensorflow.python.util import deprecation
deprecation._PRINT_DEPRECATION_WARNINGS = False
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
from datetime import datetime
from os.path import join
# from blocks import dense, conv2d, batch_norm, flatten
# from reader import read_seeds
#               ||
#               \/
from N_HANS.N_HANS___Source_Separation import blocks, reader
dense, conv2d, batch_norm, flatten = blocks.dense, blocks.conv2d, blocks.batch_norm, blocks.flatten
read_seeds = reader.read_seeds


# General flags
"""
During Training     : eval_before_train => False, eval_after_train => True
During Test         : eval_before_train => True, eval_after_train => False
restore_path        : the path to restore a trained model
checkpoints         : checkpoints directory
summaries           : summaries directory
dump_results        : directory of intermediate output of models
eval_every          : evaluation every "eval_every" batches
train_monitor_every : monitor training process every "train_monitor_every" batches
"""
FLAGS = tf.compat.v1.flags.FLAGS
tf.compat.v1.flags.DEFINE_boolean('eval_before_training', False, '')
tf.compat.v1.flags.DEFINE_boolean('eval_after_training', True, '')
tf.compat.v1.flags.DEFINE_integer('checkpoints_to_keep', 1000000, '')
tf.compat.v1.flags.DEFINE_string('restore_path', '', '')
tf.compat.v1.flags.DEFINE_string('model_name', 'nhans', '')
tf.compat.v1.flags.DEFINE_string('checkpoint_dir', './checkpoints', '')
tf.compat.v1.flags.DEFINE_string('summaries_dir', './summaries', '')
tf.compat.v1.flags.DEFINE_string('dump_results', './dump', '')
tf.compat.v1.flags.DEFINE_integer('eval_every', 3, '')  # 5000
tf.compat.v1.flags.DEFINE_integer('train_monitor_every', 3, '')  # 1000

# Training and model flags
"""
batches       : maximum batches for training
alg           : optimiser
lr            : learning rate
mom           : momentum for optimiser
w_std, b_init : parameters initialisation
bn_decay      : batch normalisation decay
train_mb      : mini-batch size for training
eval_mb       : mini-batch size for valid or test
"""
tf.compat.v1.flags.DEFINE_integer('batches', 3000000, '')
tf.compat.v1.flags.DEFINE_string('alg', 'sgd', '')
tf.compat.v1.flags.DEFINE_float('lr', 0.001, '')
tf.compat.v1.flags.DEFINE_float('mom', 0.0, '')
tf.compat.v1.flags.DEFINE_float('w_std', 0.01, '')
tf.compat.v1.flags.DEFINE_float('b_init', 0.0, '')
tf.compat.v1.flags.DEFINE_float('bn_decay', 0.95, '')
tf.compat.v1.flags.DEFINE_integer('train_mb', 64, '')  # 64
tf.compat.v1.flags.DEFINE_integer('eval_mb', 100, '')  # 100

"""
Preparation of Data Reader
"""

global treader, tsess, tph, modelname, isess, tin


def get_train_reader():
    return read_seeds('train', 40000, 6666, 16)  # 10, 1, 2


def get_eval_readers():
    return [read_seeds('valid', 5000, 0, 1)]  # 10, 0, 1


"""
N-HANS Separator Model
"""


def model(inputs, istrain):
    clean, mixed, mixedph, noisecontext, cleancontext, location, cleanpath, noisepath, snr = inputs
    nfeat = clean.shape[2].value

    def noise_resnet_block(inputs, kernel_size, stride, n_fmaps, scope_name):
        """
        Embedding network
        :param inputs : context input
        :return out   : embedding vector represents the context information
        """
        # The transformation path
        path1 = conv2d(inputs, kernel_size, [1] + stride + [1], n_fmaps, FLAGS.w_std, FLAGS.b_init, False, 'SAME',
                       scope_name + '_conv1')
        path1 = batch_norm(istrain, path1, scope_name + '_conv1')
        path1 = tf.nn.relu(path1)
        path1 = conv2d(path1, kernel_size, [1, 1, 1, 1], n_fmaps, FLAGS.w_std, FLAGS.b_init, True, 'SAME',
                       scope_name + '_conv2')

        # The identity path
        n_input_channels = inputs.shape.as_list()[3]
        if n_input_channels == n_fmaps:
            path2 = inputs
        else:
            path2 = conv2d(inputs, [1, 1], [1] + stride + [1], n_fmaps, FLAGS.w_std, FLAGS.b_init, True, 'SAME',
                           scope_name + '_transform')

        # Add and return
        assert path1.shape.as_list() == path2.shape.as_list()
        out = path1 + path2
        out = batch_norm(istrain, out, scope_name + '_addition')
        out = tf.nn.relu(out)
        return out

    def resnet_block(inputs, noiseemb, cleanemb, kernel_size, stride, n_fmaps, scope_name):
        """
        Residual block to process noisy signals, with injection of embedding vectors
        :param inputs      : input feature maps
        :param noiseposemb : positive embedding vector
        :param noisenegemb : negative embedding vector
        :param n_fmaps     : number of channels
        :return out        : output feature maps
        """

        def cont_embed(n, out_dim, scope_name):
            out = tf.constant(list(range(0, n)), dtype=tf.float32)  # [n]
            out = tf.reshape(out, [n, 1])  # [n, 1]
            out = dense(out, 50, FLAGS.w_std, 0.0, False, scope_name + '_dense1')  # [n, 50]
            out = batch_norm(istrain, out, scope_name + scope_name + '_dense1')
            out = tf.nn.relu(out)
            out = dense(out, 50, FLAGS.w_std, 0.0, False, scope_name + '_dense2')  # [n, 50]
            out = batch_norm(istrain, out, scope_name + scope_name + '_dense2')
            out = tf.nn.relu(out)
            out = dense(out, out_dim, 0.0, 0.0, False, scope_name + '_dense3')  # [n, out_dim]
            return out

        def process_noise_t_f(match_to, scope_name):
            n_fmaps = match_to.shape[3].value
            # Project the noise to fit the conv
            noise_proj = dense(noiseemb, n_fmaps, 0.0, 0.0, True, scope_name + '_noise_emb')  # [mb, n_fmaps]
            noise_proj = tf.expand_dims(noise_proj, 1)
            noise_proj = tf.expand_dims(noise_proj, 1)  # [mb, 1, 1, n_fmaps]

            clean_proj = dense(cleanemb, n_fmaps, 0.0, 0.0, True, scope_name + '_clean_emb')  # [mb, n_fmaps]
            clean_proj = tf.expand_dims(clean_proj, 1)
            clean_proj = tf.expand_dims(clean_proj, 1)  # [mb, 1, 1, n_fmaps]

            # Get the time and frequency embedding
            ts, fs = match_to.shape[1].value, match_to.shape[2].value
            tout = cont_embed(ts, n_fmaps, scope_name + '_temb')  # [ts, n_fmaps]
            tout = tf.expand_dims(tout, 1)
            tout = tf.expand_dims(tout, 0)  # [1, time, 1, n_fmaps]
            fout = cont_embed(fs, n_fmaps, scope_name + '_femb')  # [fs, n_fmaps]
            fout = tf.expand_dims(fout, 0)
            fout = tf.expand_dims(fout, 0)  # [1, 1, freq, n_fmaps]

            return noise_proj, clean_proj, tout, fout

        # The transformation path
        path1 = conv2d(inputs, [kernel_size, kernel_size], [1, stride, stride, 1],
                       n_fmaps, FLAGS.w_std, FLAGS.b_init, False,
                       'SAME', scope_name + '_conv1')  # [mb, time, freq, n_fmaps]
        noise_proj1, clean_proj1, tout1, fout1 = process_noise_t_f(path1, scope_name + '_conv1')
        path1 = path1 + noise_proj1 + clean_proj1 + tout1 + fout1
        path1 = batch_norm(istrain, path1, scope_name + '_conv1')
        path1 = tf.nn.relu(path1)
        path1 = conv2d(path1, [kernel_size, kernel_size], [1, 1, 1, 1], n_fmaps, FLAGS.w_std, FLAGS.b_init, True,
                       'SAME', scope_name + '_conv2')
        noise_proj2, clean_proj2, tout2, fout2 = process_noise_t_f(path1, scope_name + '_conv2')
        path1 = path1 + noise_proj2 + clean_proj2 + tout2 + fout2

        # The identity path
        n_input_channels = inputs.shape.as_list()[3]
        if n_input_channels == n_fmaps:
            path2 = inputs
        else:
            path2 = conv2d(inputs, [1, 1], [1, stride, stride, 1], n_fmaps, FLAGS.w_std, FLAGS.b_init, True, 'SAME',
                           scope_name + '_transform')

        # Add and return
        assert path1.shape.as_list() == path2.shape.as_list()
        out = path1 + path2
        out = batch_norm(istrain, out, scope_name + '_addition')
        out = tf.nn.relu(out)
        return out

    # The interference speaker embedding
    with tf.compat.v1.variable_scope('embedding'):
        nout = noisecontext  # [mb, noise frames, 201]
        nout = tf.expand_dims(nout, 3)
        nout = noise_resnet_block(nout, [8, 4], [3, 2], 64, 'noise_resblock1_1')  # [mb, noise frames, 201, 64]
        nout = noise_resnet_block(nout, [8, 4], [3, 2], 128, 'noise_resblock2_1')  # [mb, noise frames / 2, 201 / 2, 64]
        nout = noise_resnet_block(nout, [4, 4], [1, 1], 256, 'noise_resblock3_1')  # [mb, noise frames / 4, 201 / 4, 64]
        nout = noise_resnet_block(nout, [4, 4], [1, 2], 512,
                                  'noise_resblock4_1')  # [mb, noise frames / 8, 201 / 8, 512]
        nout = tf.nn.avg_pool2d(nout, [1, nout.shape[1].value, nout.shape[2].value, 1], [1, 1, 1, 1],
                                'VALID')  # [mb, 1, 1, 512]
        assert nout.shape.as_list()[1:3] == [1, 1]
        noiseemb = nout[:, 0, 0, :]  # [mb, 512]

    # The target speaker embedding
    with tf.compat.v1.variable_scope('embedding', reuse=True):
        nout = cleancontext  # [mb, noise frames, 201]
        nout = tf.expand_dims(nout, 3)
        nout = noise_resnet_block(nout, [8, 4], [3, 2], 64, 'noise_resblock1_1')  # [mb, noise frames, 201, 64]
        nout = noise_resnet_block(nout, [8, 4], [3, 2], 128, 'noise_resblock2_1')  # [mb, noise frames / 2, 201 / 2, 64]
        nout = noise_resnet_block(nout, [4, 4], [1, 1], 256, 'noise_resblock3_1')  # [mb, noise frames / 4, 201 / 4, 64]
        nout = noise_resnet_block(nout, [4, 4], [1, 2], 512,
                                  'noise_resblock4_1')  # [mb, noise frames / 8, 201 / 8, 512]
        nout = tf.nn.avg_pool2d(nout, [1, nout.shape[1].value, nout.shape[2].value, 1], [1, 1, 1, 1],
                                'VALID')  # [mb, 1, 1, 512]
        assert nout.shape.as_list()[1:3] == [1, 1]
        cleanemb = nout[:, 0, 0, :]  # [mb, 512]

    # Processing the mixed signal
    out = mixed  # [mb, context frames, 201]
    out = tf.expand_dims(out, 3)
    out = resnet_block(out, noiseemb, cleanemb, 4, 1, 64, 'resblock1_1')
    out = resnet_block(out, noiseemb, cleanemb, 4, 1, 64, 'resblock1_2')
    out = resnet_block(out, noiseemb, cleanemb, 4, 2, 128, 'resblock2_1')
    out = resnet_block(out, noiseemb, cleanemb, 4, 1, 128, 'resblock2_2')
    out = resnet_block(out, noiseemb, cleanemb, 3, 2, 256, 'resblock3_1')
    out = resnet_block(out, noiseemb, cleanemb, 3, 1, 256, 'resblock3_2')
    out = resnet_block(out, noiseemb, cleanemb, 3, 2, 512, 'resblock4_1')
    out = resnet_block(out, noiseemb, cleanemb, 3, 1, 512, 'resblock4_2')  # [mb, context frames / 8, 201 / 8, 512]

    # final layers
    out = conv2d(out, [out.shape[1].value, 1], [1, 1, 1, 1],
                 512, FLAGS.w_std, FLAGS.b_init, False,
                 'VALID', 'last_conv')  # [mb, 1, 201 / 8, 512]
    out = batch_norm(istrain, out, 'last_conv')
    out = tf.nn.relu(out)
    out = flatten(out)  # [mb,  (201 / 8) * 512]
    out = dense(out, nfeat, 0.0, 0.0, True, 'last_dense')  # [mb, 201]
    mixed_central = mixed[:, FLAGS.window_frames // 2, :]  # [mb, 201]
    denoised = mixed_central + out  # [mb, 201]

    # Loss
    se = tf.square(denoised - clean[:, 0, :])  # [mb, 201]
    imp_factor = np.linspace(2, 1, nfeat, dtype=np.float32).reshape((1, nfeat))
    example_loss = tf.reduce_mean(se * tf.constant(imp_factor), axis=1)
    loss = tf.reduce_mean(example_loss)

    monitors = {'loss': loss}
    outputs = {'loss': example_loss, 'mixed': mixed_central, 'denoised': denoised, 'clean': clean[:, 0, :],
               'mixedph': mixedph[:, 0, :], 'location': location, 'cleanpath': cleanpath,
               'noisepath': noisepath, 'snr': snr}
    return loss, monitors, outputs


"""
Restoration from Estimated Spectrum
"""


def evaluate(outputs, ereader, step):
    """
    Restoration from Estimated Spectrum
    :param outputs : estimated spectrograms
    :param ereader : evaluation data reader
    :param step    : global step
    """
    print(ereader.name)

    # Print the loss
    loss = outputs['loss'].mean()
    print('loss: {}'.format(loss))

    # Create a summary
    summary = tf.compat.v1.Summary()
    summary.value.add(tag='{}_{}'.format(ereader.name, 'loss'), simple_value=loss)

    # Construct the tf graph for reconstruction
    g = tf.Graph()
    with g.as_default():
        stft = tf.compat.v1.placeholder(dtype=tf.complex64, shape=[None, 201])
        samplestf = tf.signal.inverse_stft(tf.expand_dims(stft, 0), ereader.frame_length, ereader.frame_step,
                                           ereader.frame_length, window_fn=tf.signal.inverse_stft_window_fn(ereader.frame_step, forward_window_fn=functools.partial(tf.signal.hann_window, periodic=True)))[0, :]
        sess = tf.compat.v1.Session()

    # Reconstruct some audio
    startlocations = np.where(outputs['location'] == 0)[0]
    for i in range(0, len(startlocations)):
        # Start and end location
        s = startlocations[i]
        if i == len(startlocations) - 1:
            e = len(outputs['mixed'])
        else:
            e = startlocations[i + 1]

        # Collect frames
        mixedmagframes = outputs['mixed'][s:e]
        mixedphframes = outputs['mixedph'][s:e]
        denoisedmagframes = outputs['denoised'][s:e]

        # Undo the log operation
        mixedmagframes = np.exp(mixedmagframes)
        denoisedmagframes = np.exp(denoisedmagframes)

        # Reconstruct complex spectrogram
        mixedspect = mixedmagframes * (np.exp(1j * mixedphframes))
        denoisedspect = denoisedmagframes * (np.exp(1j * mixedphframes))

        # Reconstruct waveform
        mixedsamples = sess.run(samplestf, feed_dict={stft: mixedspect})
        denoisedsamples = sess.run(samplestf, feed_dict={stft: denoisedspect})

        # Save the waveforms
        cleanpath, noisepath, snr = outputs['cleanpath'][s], outputs['noisepath'][s], outputs['snr'][s]
        cleanpath, noisepath = cleanpath.decode('utf-8'), noisepath.decode('utf-8')
        cleanpath = cleanpath.split('/')[-1][:-4]
        noisepath = noisepath.split('/')[-1][:-4]
        mixed_filename = '{}_{}_{}_{}_{}_{}.wav'.format(modelname, step, cleanpath, noisepath, snr, 'mixed')
        denoised_filename = '{}_{}_{}_{}_{}_{}.wav'.format(modelname, step, cleanpath, noisepath, snr, 'denoised')
        scipy.io.wavfile.write(os.path.join(FLAGS.wav_dump_folder, mixed_filename), 16000, mixedsamples)
        scipy.io.wavfile.write(os.path.join(FLAGS.wav_dump_folder, denoised_filename), 16000, denoisedsamples)


"""
Main Loop for Training and Evaluation with Monitors
"""


def main_loop():
    global treader, tsess, tph, modelname, isess, tin

    # Print FLAGS
    _ = FLAGS.lr
    print('----------------------------- FLAGS VALUES --------------------------------')
    for k in sorted(FLAGS.__flags.keys()):
        print('{}: {}'.format(k, getattr(FLAGS, k)))
    print('')

    # Print General model information
    print('----------------------- DATA LOADING, MODEL PREPARING -------------------------')
    time_stamp = str(datetime.now()).replace(':', '-').replace(' ', '_')
    modelname = FLAGS.model_name if FLAGS.model_name else str(time_stamp)
    print('model_name: {}'.format(modelname))

    ig = tf.Graph()
    with ig.as_default():
        with tf.device('/cpu:0'):
            # Preparing the reader
            treader = get_train_reader()
            treader.preparations()

            # Readers
            ereaders = get_eval_readers()

        isess = tf.compat.v1.Session(graph=ig, config=tf.compat.v1.ConfigProto(log_device_placement=False,
                                                                               allow_soft_placement=True))

    # Training model settings
    tg = tf.Graph()
    with tg.as_default():
        # Session
        tsess = tf.compat.v1.Session(graph=tg, config=tf.compat.v1.ConfigProto(log_device_placement=False,
                                                                               allow_soft_placement=True))

        # The model
        with tf.device('/gpu:0'):
            tin = treader.get_inputs()
            loss, monitors, _ = model(tin, True)
            # The global step
            global_step = tf.Variable(0, trainable=False)
            # Optimization
            if FLAGS.alg == 'sgd':
                opt = tf.compat.v1.train.GradientDescentOptimizer(FLAGS.lr)
            elif FLAGS.alg == 'momentum':
                opt = tf.train.MomentumOptimizer(FLAGS.lr, FLAGS.mom)
            elif FLAGS.alg == 'rmsprop':
                opt = tf.train.RMSPropOptimizer(FLAGS.lr, momentum=FLAGS.mom)
            elif FLAGS.alg == 'adadelta':
                opt = tf.train.AdadeltaOptimizer(FLAGS.lr)
            elif FLAGS.alg == 'adagrad':
                opt = tf.train.AdagradOptimizer(FLAGS.lr)
            elif FLAGS.alg == 'adam':
                opt = tf.train.AdamOptimizer(learning_rate=FLAGS.lr)
            grads_and_vars = opt.compute_gradients(loss, tf.compat.v1.trainable_variables())
            train_op = opt.apply_gradients(grads_and_vars)
            print('#trainable variables: {}'.format(
                sum([x.get_shape().num_elements() for x in tf.compat.v1.trainable_variables()])))
            print('#non-trainable variables: {}'.format(sum([x.get_shape().num_elements() for x in
                                                             set(tf.compat.v1.global_variables()) - set(
                                                                 tf.compat.v1.trainable_variables())])))
            assert not tf.compat.v1.local_variables()  # Should be empty always, otherwise change the lines above

        # Save, restore and init
        tsaver = tf.compat.v1.train.Saver(tf.compat.v1.global_variables(), max_to_keep=FLAGS.checkpoints_to_keep)
        if FLAGS.restore_path:
            # Restore the variables values
            print('Restoring model from {}'.format(FLAGS.restore_path))
            restore_saver = tf.compat.v1.train.Saver(tf.compat.v1.global_variables())
            restore_saver.restore(tsess, FLAGS.restore_path)
        else:
            tsess.run(tf.compat.v1.global_variables_initializer())

    # Evaluation model settings
    eg = tf.Graph()
    with eg.as_default():
        esess = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(allow_soft_placement=True))
        with tf.device('/gpu:0'):
            ein = ereaders[0].get_inputs()
            _, _, outputs = model(ein, False)
        esaver = tf.compat.v1.train.Saver(tf.compat.v1.global_variables())

    # Training loop
    print('--------------------------------- TRAINING! ------------------------------------')

    # Init aggregators and names
    aggregators = {k: np.zeros(shape=monitors[k].shape.as_list()) for k in monitors.keys()}

    # Operations that need to be executed in every training step
    step_ops = {'train_op': train_op}
    step_ops.update(monitors)

    # Measure time
    start = time.time()

    # The step
    tstep = 0

    try:
        # Start queue runners
        with ig.as_default():
            with tf.device('/cpu:0'):
                isess.run(tf.compat.v1.local_variables_initializer())
                icoord = tf.train.Coordinator()
                ithreads = tf.train.start_queue_runners(sess=isess, coord=icoord, start=False)
                queue_stats = tf.stack(
                    [x.queue.size() for x in tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.QUEUE_RUNNERS)])
                for t in ithreads:
                    t.start()

        def save_and_eval():
            # Saving and restoring
            print('Saving and restoring the model')
            tsaver.save(tsess,
                        os.path.join(FLAGS.checkpoint_dir, modelname),
                        global_step=tstep,
                        latest_filename=modelname)
            ckpt = tf.train.get_checkpoint_state(FLAGS.checkpoint_dir, latest_filename=modelname)
            if ckpt and ckpt.model_checkpoint_path:
                print('Restoring from file: {}'.format(ckpt.model_checkpoint_path))
                global_step = ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]
            else:
                print('No checkpoint file found')
                return
            esaver.restore(esess, ckpt.model_checkpoint_path)

            # Print the header
            print('----------------- TEST MONITOR ----------------------')

            for ereader in ereaders:
                # Initialize aggregators
                aggregators = {k: [] for k in outputs.keys()}

                # Build the graph again
                try:
                    with ig.as_default():
                        with tf.device('/cpu:0'):
                            ereader.preparations()
                            isess.run(tf.compat.v1.local_variables_initializer())
                            nowcoord = tf.train.Coordinator()
                            nowthreads = []
                            qr1 = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.QUEUE_RUNNERS)[-3]
                            qr2 = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.QUEUE_RUNNERS)[-2]
                            qr3 = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.QUEUE_RUNNERS)[-1]
                            nowthreads.extend(qr1.create_threads(isess, coord=nowcoord, daemon=True, start=True))
                            nowthreads.extend(qr2.create_threads(isess, coord=nowcoord, daemon=True, start=True))
                            nowthreads.extend(qr3.create_threads(isess, coord=nowcoord, daemon=True, start=True))

                    i = 0
                    while True:
                        i += 1

                        input_feed_dict = {ereader.mb: FLAGS.eval_mb}
                        try:
                            batch = isess.run(ereader.batch, feed_dict=input_feed_dict)
                        except tf.errors.OutOfRangeError:
                            break
                        QUICK_MODE = True
                        if QUICK_MODE and i == 2:
                            break
                        model_feed_dict = {a: b for (a, b) in zip(ein, batch)}
                        results = esess.run(outputs, model_feed_dict)

                        # concatenation aggregation of certain channels
                        for k in results.keys():
                            aggregators[k].append(results[k])
                finally:
                    nowcoord.request_stop()
                    nowcoord.join(nowthreads)

                # Concatenate results over different batches
                for k in results.keys():
                    aggregators[k] = np.concatenate(aggregators[k])

                # Dump activations
                if FLAGS.dump_results:
                    for k in results.keys():
                        np.save(join(FLAGS.dump_results, '{}_{}_{}_{}'.format(modelname, ereader.name, tstep, k)),
                                aggregators[k])

                # Evaluate
                evaluate(aggregators, ereader, global_step)

            # Print the bottom part
            print('-----------------------------------------------------')
            print('')

        # Eval before training
        if FLAGS.eval_before_training:
            save_and_eval()

        # Training loop
        while tstep < FLAGS.batches and not icoord.should_stop():
            # Running the step
            input_feed_dict = {treader.mb: FLAGS.train_mb}
            # t1 = time.time()
            batch = isess.run(treader.batch, feed_dict=input_feed_dict)
            # t2 = time.time()
            model_feed_dict = {a: b for (a, b) in zip(tin, batch)}
            results = tsess.run(step_ops, model_feed_dict)

            # Aggregate the train monitoring channels
            for key in monitors.keys():
                aggregators[key] += results[key]

            # Increase the step. 'step' always contains the number of gradient updates that were performed.
            tstep += 1

            # Train monitoring
            if tstep % FLAGS.train_monitor_every == 0:
                # Do some printing
                print('----- TRAIN MONITOR AFTER ANOTHER {} BATCHES ------------'.format(FLAGS.train_monitor_every))
                print('step number: {}'.format(tstep))
                for k in sorted(aggregators.keys()):
                    aggregators[k] = aggregators[k] / float(FLAGS.train_monitor_every)
                    print('{}: {}'.format(k, aggregators[k]))  # loss
                end = time.time()
                print('queues stats: {}'.format(isess.run(queue_stats)))
                print('seconds elapsed: {}'.format(end - start))
                start = end
                # Add summaries and reset aggregators
                summary = tf.compat.v1.Summary()
                summary.value.add(tag='{}_{}'.format(treader.name, 'loss'), simple_value=aggregators['loss'])
                # treader.swriter.add_summary(summary, tstep)
                aggregators = {k: 0.0 for k in monitors.keys()}
                print('---------------------------------------------------------')
                print('')

            # Save and eval
            if tstep % FLAGS.eval_every == 0:
                save_and_eval()

        # Eval after training
        if FLAGS.eval_after_training:
            save_and_eval()

    finally:
        icoord.request_stop()
        time.sleep(1)
        isess.close()
        tsess.close()
        esess.close()
        icoord.join(ithreads)
  
      
if __name__ == '__main__':
    main_loop()
