########################################################################################################################
#                                          N-HANS speech separator: reader                                             #
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
#   Description:      Data reader pipeline.                                                                            #
#   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:      June. 17, 2020                                                                                   #
#   Dependence Files: xxx                                                                                              #
#   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 numpy as np
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import tensorflow as tf
from tensorflow.python.util import deprecation
deprecation._PRINT_DEPRECATION_WARNINGS = False
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
import random
import time
import hashlib
import pickle as p
from scipy.io.wavfile import read as wavread


"""
context_frames  : number of context frames          (reference signal)
window_frames   : number of frames of noisy signal  (input signal) 
random_slices   : number of slices (each slice contains one input signal and reference signals) 
eval_seeds      : 'valid' or 'test'. In training, evaluation is applied for 'valid' dataset, and in test, for 'test' dataset
wav_dump_folder : the folder to save denoised signals
"""
FLAGS = tf.compat.v1.flags.FLAGS
tf.compat.v1.flags.DEFINE_integer('context_frames', 200, '')
tf.compat.v1.flags.DEFINE_integer('window_frames', 35, '')
tf.compat.v1.flags.DEFINE_integer('random_slices', 50, '')
tf.compat.v1.flags.DEFINE_string('eval_seeds', 'valid', '')
tf.compat.v1.flags.DEFINE_string('wav_dump_folder', './wav_dump/', '')
tf.compat.v1.flags.DEFINE_string('speech_wav_dir', './speech_wav_dir/', '')


"""
Create Pickle List for Training, Evaluation and Test
"""


def create_seeds():
    """
    Create 3 pickle files: train.pkl, valid.pkl and test.pkl for speech audio files
    :param speech_wav_dir: the directory contains three subfolders, train, valid and test, each includes audio .wav files
    """
    speech_wav_dir = FLAGS.speech_wav_dir
    train_folder_clean = speech_wav_dir + 'train'
    valid_folder_clean = speech_wav_dir + 'valid'
    test_folder_clean = speech_wav_dir + 'test'
    trainseeds = []
    validseeds = []
    testseeds = []
    wavfolders = [train_folder_clean, valid_folder_clean, test_folder_clean]
    seedsfolders = [trainseeds, validseeds, testseeds]

    for ii in range(3):
        wavfolder = wavfolders[ii]
        seedfolder = seedsfolders[ii]
        for root, directories, filenames in os.walk(wavfolder):
            for filename in filenames:
                if filename.endswith('.wav'):
                    seedfolder.append(os.path.join(root, filename))

    train_pkl = speech_wav_dir + 'train.pkl'
    valid_pkl = speech_wav_dir + 'valid.pkl'
    test_pkl = speech_wav_dir + 'test.pkl'
    p.dump(trainseeds, open(train_pkl, "wb"))
    p.dump(validseeds, open(valid_pkl, "wb"))
    p.dump(testseeds, open(test_pkl, "wb"))


"""
Additional Functions for Generating Audio Samples
combine_signals: generating training, valid and test .wav files
"""


def read_wav(in_path):
    rate, samples = wavread(in_path)
    assert rate == 16000
    assert samples.dtype == 'int16'
    if len(samples.shape) > 1:
        samples = samples.mean(axis=1)
    assert len(samples.shape) == 1
    return samples


def domixing(cleansamples, noisesamples, snr):
    # Repeat the noise if it is shorter than the speech, or shorter it if it's longer
    nse = noisesamples
    while len(cleansamples) - len(nse) > 0:  # Make noise longer
        diff = len(cleansamples) - len(nse)
        nse = np.concatenate([nse, noisesamples[:diff]], axis=0)
    if len(cleansamples) - len(noisesamples) < 0:  # Make noise shorter
        nse = noisesamples[:len(cleansamples)]
    sig = cleansamples

    # Power of signal and noise
    psignal = sum(abs(sig) * abs(sig)) / sig.shape[0]
    pnoise = sum(abs(nse) * abs(nse)) / nse.shape[0]

    # Compute scale factor
    if pnoise == 0:
        K = 1
    else:
        K = (psignal / pnoise) * pow(10, -snr / 10.0)
    K = np.sqrt(K)

    # Mix
    noise_scaled = K * nse  # Scale the noise
    mixed = sig + noise_scaled  # Mix
    mixed = mixed / (max(abs(mixed)) + 0.000001)  # Normalize
    return mixed, K


def combine_signals(istrain, cleanpath, noisepath):
    try:
        # Read Wavs
        cleansamples = read_wav(cleanpath)
        noisesamples = read_wav(noisepath)

        # Normalize
        cleansamples = cleansamples / (max(abs(cleansamples)) + 0.000001)
        noisesamples = noisesamples / (max(abs(noisesamples)) + 0.000001)
        cleansamples = cleansamples.astype(np.float32)
        noisesamples = noisesamples.astype(np.float32)

        # Cut the end to have an exact number of frames
        cleansamples = cleansamples[:-((len(cleansamples) - 400) % 160)]

        # Choose SNR
        SNRs = [-5, -3, -1, 0, 1, 3, 5]

        if istrain:
            snr = SNRs[random.randint(0, len(SNRs) - 1)]
            mixed, K = domixing(cleansamples, noisesamples, snr)
            return cleansamples, noisesamples * K, mixed, np.array(snr, dtype=np.int32)

        else:
            # The noise is the beginning of the file, and mix the rest
            snrid = int(hashlib.md5(cleanpath).hexdigest()[:8], 16) % len(SNRs)
            snr = SNRs[snrid]
            mixed, K = domixing(cleansamples, noisesamples, snr)
            return cleansamples, noisesamples * K, mixed, np.array(snr, dtype=np.int32)
    except:
        print('error in threads')
        print(cleanpath, noisepath)


"""
Global FLAGS and Training FLAGS
"""


class read_seeds:
    """
    name              : 'train', 'valid' or 'test'. To generate data reader.
    queuesize         : the size of data queue
    min_after_dequeue : minimum seeds left in queue after dequeue operation
    nthreads          : number of threads processing dequeue
    """

    def __init__(self, name, queuesize, min_after_dequeue, nthreads):
        self.Fs = 16000
        self.frame_length = int(self.Fs * 0.025)
        self.frame_step = int(self.Fs * 0.010)
        self.queuesize = queuesize
        self.min_after_dequeue = min_after_dequeue
        self.window_frames = FLAGS.window_frames
        self.random_slices = FLAGS.random_slices
        self.context_frames = FLAGS.context_frames
        self.eval_stride = 1
        self.istrain = name == 'train'
        self.nthreads = nthreads
        self.name = name

        if self.istrain != True:
            name = FLAGS.eval_seeds

        speech_wav_dir = FLAGS.speech_wav_dir
        if name == 'train':
            self.seedspath = speech_wav_dir + 'train.pkl'
        elif name == 'valid':
            self.seedspath = speech_wav_dir + 'valid.pkl'
        elif name == 'test':
            self.seedspath = speech_wav_dir + 'test.pkl'

    def preparations(self):
        """
        Preparation of the data queue.
        """
        self.seeds = []
        with open(self.seedspath, 'rb') as f:
            seedlist = p.load(f)
            seedlist = tf.constant(np.array(seedlist))
            self.seeds.append(seedlist)

        # Construct seeds queues
        self.seedsqs = []
        for i, seed in enumerate(self.seeds):
            num_epochs = None if self.istrain else 1
            seedsq = tf.train.input_producer(seed, capacity=50000, shuffle=self.istrain, num_epochs=num_epochs)
            self.epochcounter = None if self.istrain else tf.compat.v1.local_variables()[-1]
            self.seedsqs.append(seedsq)

        # Get the examples from seeds. It is a list of tensors, with first dimension being the example.
        examples = self.get_examples(self.istrain, self.seedsqs)

        # Create examples queue and enqueue the examples
        dtypes = [x.dtype for x in examples]
        shapes = [x.shape.as_list()[1:] for x in examples]
        if self.istrain:
            exq = tf.queue.RandomShuffleQueue(capacity=self.queuesize, min_after_dequeue=self.min_after_dequeue,
                                              dtypes=dtypes, shapes=shapes)
        else:
            exq = tf.queue.FIFOQueue(capacity=self.queuesize, dtypes=dtypes, shapes=shapes)
        enqop = exq.enqueue_many(examples)

        # The dequeue operation and the mb placeholder
        mb = tf.compat.v1.placeholder(dtype=tf.int32, shape=[], name=self.name + '_mb')
        if self.istrain:
            batch = exq.dequeue_many(mb)
        else:
            batch = exq.dequeue_up_to(mb)

        # Create queue runner
        qr = tf.train.QueueRunner(exq, [enqop] * self.nthreads)
        tf.train.add_queue_runner(qr, tf.compat.v1.GraphKeys.QUEUE_RUNNERS)

        # Set some members
        self.exq = exq
        self.batch = batch
        self.mb = mb

    def get_examples(self, istrain, seedsqs):
        """
        get model inputs from audio samples
        :param istrain: boolean value. If True, the examples are randomly trauncated from audio samples; If False, the examples are trauncated frame by frame from audio samples
        :param seedsqs: seed list of speech files
        :return:
        """
        target_seed = seedsqs[0].dequeue()
        noise_seed = seedsqs[0].dequeue()
        target_wav, noise_wav, mix_wav, snr = tf.py_func(combine_signals, [istrain, target_seed, noise_seed],
                                                         [tf.float32, tf.float32, tf.float32, tf.int32])

        # flatten
        target_wav, noise_wav, mix_wav = [tf.reshape(x, [-1]) for x in (target_wav, noise_wav, mix_wav)]

        # stft
        target_fft = tf.signal.stft(target_wav, self.frame_length, self.frame_step,
                                            fft_length=self.frame_length)
        noise_fft = tf.signal.stft(noise_wav, self.frame_length, self.frame_step, fft_length=self.frame_length)
        mix_fft = tf.signal.stft(mix_wav, self.frame_length, self.frame_step, fft_length=self.frame_length)

        # Magnitude
        target_mag, noise_mag, mix_mag = [tf.abs(x) for x in [target_fft, noise_fft, mix_fft]]
        # Phase
        mix_phase = tf.math.angle(mix_fft)
        # log
        target_log, noise_log, mix_log = [tf.math.log(x + 1e-5) for x in [target_mag, noise_mag, mix_mag]]

        if istrain:
            # padding
            target_log, noise_log, mix_log, mix_phase = [self.pad_1D_for_windowing(x, self.window_frames) for x in
                                                         [target_log, noise_log, mix_log, mix_phase]]

            # Get n random window-label pairs
            targetslices = []
            mixedslices = []
            mixedphslices = []
            noise_context_slices = []
            clean_context_slices = []
            for _ in range(self.random_slices):
                # Crop a window
                cleancropped, mixedcropped, mixedphcropped, noisecontext, cleancontext = \
                    self.synchronized_1D_crop(target_log, noise_log, mix_log, mix_phase, self.window_frames,
                                              self.context_frames)

                # Add to list
                for (lst, tnsr) in zip(
                        [targetslices, mixedslices, mixedphslices, noise_context_slices, clean_context_slices],
                        [cleancropped, mixedcropped, mixedphcropped, noisecontext, cleancontext]):
                    lst.append(tnsr)

            # Stack and return
            return tf.stack(targetslices, axis=0), tf.stack(mixedslices, axis=0), tf.stack(mixedphslices, axis=0), \
                   tf.stack(noise_context_slices, axis=0), tf.stack(clean_context_slices, axis=0), \
                   tf.constant([0] * self.random_slices, dtype=tf.int32), tf.constant([''] * self.random_slices,
                                                                                      dtype=tf.string), \
                   tf.constant([''] * self.random_slices, dtype=tf.string), tf.constant([0] * self.random_slices,
                                                                                        dtype=tf.int32)
        else:
            targetslices = self.strided_crop(target_log[self.context_frames:], 1, self.eval_stride)
            mixedslices = self.strided_crop(mix_log[self.context_frames:], self.window_frames, self.eval_stride)
            mixedphslices = self.strided_crop(mix_phase[self.context_frames:], 1, self.eval_stride)
            noisecontext = noise_log[:self.context_frames]
            noisecontext = tf.reshape(noisecontext, [self.context_frames, noisecontext.shape[1].value])
            noise_context_slices = tf.tile(tf.expand_dims(noisecontext, 0), [tf.shape(targetslices)[0], 1, 1])
            cleancontext = target_log[:self.context_frames]
            cleancontext = tf.reshape(cleancontext, [self.context_frames, cleancontext.shape[1].value])
            clean_context_slices = tf.tile(tf.expand_dims(cleancontext, 0), [tf.shape(targetslices)[0], 1, 1])
            locations = tf.range(0, tf.shape(targetslices)[0], dtype=tf.int32)
            target_seeds = tf.tile([target_seed], [tf.shape(targetslices)[0]])
            noise_seeds = tf.tile([noise_seed], [tf.shape(targetslices)[0]])
            snrs = tf.tile([snr], [tf.shape(targetslices)[0]])
            return targetslices, mixedslices, mixedphslices, noise_context_slices, clean_context_slices, locations, target_seeds, noise_seeds, snrs

    def pad_1D_for_windowing(self, tensor, length):
        """
        Padding the input spectrogramm
        :param tensor: shape of [Times x Frequencies]
        :param length: equals to window frames
        :return: padded spectrogram
        """
        len_before = ((length + 1) // 2) - 1
        len_after = length // 2
        return tf.pad(tensor, [[len_before, len_after], [0, 0]])

    def synchronized_1D_crop(self, clean, noise, mixed, mixedph, winlength, contextlength):
        winmaxval = tf.shape(mixed)[0] - winlength
        winstart = tf.random.uniform([], dtype=tf.int32, minval=0, maxval=winmaxval + 1)

        # Crop window
        mixedcropped = mixed[winstart:winstart + winlength]
        cleancropped = clean[winstart:winstart + winlength]
        mixedphcropped = mixedph[winstart:winstart + winlength]
        mixedcropped = tf.reshape(mixedcropped, [winlength, mixedcropped.shape[1].value])
        cleancropped = tf.reshape(cleancropped, [winlength, cleancropped.shape[1].value])
        mixedphcropped = tf.reshape(mixedphcropped, [winlength, mixedphcropped.shape[1].value])

        # Get the center of the clean and mixed phase
        cleancropped = cleancropped[self.window_frames // 2]
        cleancropped = tf.expand_dims(cleancropped, 0)
        mixedphcropped = mixedphcropped[self.window_frames // 2]
        mixedphcropped = tf.expand_dims(mixedphcropped, 0)

        # The rest of the noise
        cleanrest = tf.concat([clean[:winstart], clean[winstart + winlength:]], axis=0)
        noiserest = tf.concat([noise[:winstart], noise[winstart + winlength:]], axis=0)

        # Get the noise part
        restmaxval = tf.shape(noiserest)[0] - contextlength
        reststart = tf.random.uniform([], dtype=tf.int32, minval=0, maxval=restmaxval + 1)
        noisecontext = noiserest[reststart:reststart + contextlength]
        noisecontext = tf.reshape(noisecontext, [contextlength, noisecontext.shape[1].value])

        restmaxval = tf.shape(cleanrest)[0] - contextlength
        reststart = tf.random.uniform([], dtype=tf.int32, minval=0, maxval=restmaxval + 1)
        cleancontext = cleanrest[reststart:reststart + contextlength]
        cleancontext = tf.reshape(cleancontext, [contextlength, cleancontext.shape[1].value])

        return cleancropped, mixedcropped, mixedphcropped, noisecontext, cleancontext

    def strided_crop(self, tensor, length, stride):
        # we assume that we have a length dimension and a feature dimension
        assert len(tensor.shape) == 2
        n_features = int(tensor.shape[1])
        padded = self.pad_1D_for_windowing(tensor, length)
        windows = tf.extract_image_patches(tf.expand_dims(tf.expand_dims(padded, axis=0), axis=3),
                                           ksizes=[1, length, n_features, 1],
                                           strides=[1, stride, n_features, 1],
                                           rates=[1, 1, 1, 1],
                                           padding='VALID')
        return tf.reshape(windows, [-1, length, n_features])

    def get_inputs(self):
        if not hasattr(self, 'ph'):
            self.ph = [tf.compat.v1.placeholder(dtype=tf.float32, shape=[None, 1, 201], name='cleanph'),
                       tf.compat.v1.placeholder(dtype=tf.float32, shape=[None, self.window_frames, 201],
                                                name='mixedph'),
                       tf.compat.v1.placeholder(dtype=tf.float32, shape=[None, 1, 201], name='mixedphaseph'),
                       tf.compat.v1.placeholder(dtype=tf.float32, shape=[None, self.context_frames, 201],
                                                name='noisecontextph'),
                       tf.compat.v1.placeholder(dtype=tf.float32, shape=[None, self.context_frames, 201],
                                                name='cleancontextph'),
                       tf.compat.v1.placeholder(dtype=tf.int32, shape=[None], name='locationph'),
                       tf.compat.v1.placeholder(dtype=tf.string, shape=[None], name='cleanpathph'),
                       tf.compat.v1.placeholder(dtype=tf.string, shape=[None], name='noisepathph'),
                       tf.compat.v1.placeholder(dtype=tf.int32, shape=[None], name='snrph')]
        return self.ph


if __name__ == '__main__':
    g = tf.Graph()
    with g.as_default():
        reader = read_seeds(name='train', queuesize=200, min_after_dequeue=100, nthreads=16)
        with tf.device('/cpu:0'):
            reader.preparations()

        print('starting queue runners')
        sess = tf.compat.v1.Session(graph=g, config=tf.ConfigProto(allow_soft_placement=True))
        sess.run(tf.compat.v1.local_variables_initializer())
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)
        queue_stats = tf.stack([x.queue.size() for x in tf.get_collection(tf.compat.v1.GraphKeys.QUEUE_RUNNERS)])
        feed_dict = {reader.mb: 64}

        with tf.device('/gpu:0'):
            clean, mixed, mixedph, noisecontext, cleancontext, location, cleanpath, noisepath, snr = reader.get_inputs()

        for i in range(1000):
            print(i)
            t1 = time.time()
            clean_, mixed_, mixedph_, noisecontext_, cleancontext_, location_, cleanpath_, noisepath_, snr_ = sess.run(
                reader.batch, feed_dict=feed_dict)
            t2 = time.time()
            print(sess.run(queue_stats), t2 - t1)
