#!/usr/bin/env python3

"""Real-time single-channel denoising

"""
from __future__ import division
from __future__ import print_function
#import argparse
#import os
#import time
import jack
try:
    import queue  # Python 3.x
except ImportError:
    import Queue as queue  # Python 2.x
import sys
import threading
import numpy as np

import tensorflow as tf
from tensorflow.contrib.layers import xavier_initializer, flatten, fully_connected
#from tensorflow.contrib.lite import toco_convert

#if sys.version_info < (3, 0):
#    import signal
#    signal.signal(signal.SIGINIT,signal.SIG_DFL)
#else:
#    pass

import keras
from keras.layers import Input, Dense, Conv1D, Conv2D, Conv2DTranspose, BatchNormalization
from keras.layers import LeakyReLU, PReLU, Reshape, Concatenate, Flatten, Add, Lambda
from keras.models import Sequential, Model, model_from_json
from keras.optimizers import Adam
from keras.callbacks import TensorBoard
keras_backend = tf.keras.backend
keras_initializers = tf.keras.initializers
from data_ops import data_preprocess, de_emph
from file_ops import get_modeldirname
#from models import *

# Define optimizers
g_opt = keras.optimizers.Adam(lr=0.0002)

modeldir = "./AECNN_15_light"
print ("Loading model from " + modeldir + "/Gmodel")
json_file = open(modeldir + "/Gmodel.json", "r")
loaded_model_json = json_file.read()
json_file.close()
G_loaded = model_from_json(loaded_model_json)
G_loaded.compile(loss='mean_squared_error', optimizer=g_opt)
G_loaded.load_weights(modeldir + "/Gmodel.h5")

buffersize = 1
if buffersize < 1:
    print('buffersize must be at least 1')

q = queue.Queue(maxsize=buffersize)
qin = queue.Queue(maxsize=buffersize)
event = threading.Event()


def print_error(*args):
    print(*args, file=sys.stderr)


def xrun(delay):
    print_error("An xrun occured, increase JACK's period size?")


def shutdown(status, reason):
    print_error('JACK shutdown!')
    print_error('status:', status)
    print_error('reason:', reason)
    event.set()


def stop_callback(msg=''):
    if msg:
        print_error(msg)
    for port in client.outports:
        port.get_array().fill(0)
    event.set()
    #raise jack.CallbackExit


#drops = 0
def process(frames):
    #global drops
    if frames != blocksize:
        stop_callback('blocksize must not be changed, I quit!')
    try:
	datain=client.inports[0].get_array()
    	qin.put(datain)
    	data = q.get_nowait()
	client.outports[0].get_array()[:] = data.T
    except queue.Empty:
	#drops += 1
        stop_callback('Buffer is empty: increase buffersize?')


try:
    client = jack.Client("thru_client")
    blocksize = client.blocksize
    samplerate = client.samplerate
    client.set_xrun_callback(xrun)
    client.set_shutdown_callback(shutdown)
    client.set_process_callback(process)

    client.inports.register('in_{0}'.format(1))
    client.outports.register('out_{0}'.format(1))
    i=client.inports[0]
    data=np.zeros(blocksize,)

    noisy=np.concatenate((np.zeros(blocksize,),data))
    noisy.shape = (1,2*blocksize) # first time run
    noisy = np.expand_dims(noisy, axis = 2)
    clean = G_loaded.predict(noisy)
    clean = np.reshape(clean, (1,2*blocksize))
    clean = np.reshape(clean, (-1,)) # make it to 1d by dropping the extra dimension
    #total = 0

    noisy=np.zeros(blocksize,)
    q.put_nowait(noisy)  # Pre-fill queue
    qin.put_nowait(noisy)

    #np.set_printoptions(precision=1, threshold=np.inf, linewidth=np.inf)
    capture = client.get_ports(is_physical=True, is_output=True)
    playback = client.get_ports(is_physical=True, is_input=True, is_audio=True)
    o=client.outports[0]
    timeout = blocksize * buffersize / samplerate
    print(timeout)

    datap=np.zeros(blocksize,)
    cleanp=np.zeros(blocksize,)
    with client:
	i.connect(capture[0])
	# Connect mono file to stereo output
	o.connect(playback[0])
	o.connect(playback[1])
	while True:
	#	total+=1
		data = qin.get_nowait()
		#t=time.time()
		# processing input
		noisy=np.concatenate((datap,data))
		datap=data
		noisy=noisy.T
		noisy.shape=(1,2*blocksize)
		#clean=np.zeros(frames,)
		#noisy = data_preprocess(noisy, preemph=0.95)
		noisy = np.expand_dims(noisy, axis = 2)
		clean = G_loaded.predict(noisy)
		#clean = noisy
		clean = np.reshape(clean, (1,2*blocksize))
		clean = np.reshape(clean, (-1,)) # make it to 1d by dropping the extra dimension
		#clean = de_emph(clean, coeff=0.95)
		#print(time.time()-t)
		dataout=(cleanp+clean[0:blocksize])/2
		cleanp=clean[blocksize:2*blocksize]
		q.put(dataout)
	event.wait()  # Wait until playback is finished
except KeyboardInterrupt:
    #print(drops/total)
    print('\nInterrupted by user')
except (queue.Full):
    # A timeout occured, i.e. there was an error in the callback
    #parser.exit(1)
    print('\nQueue full')
    KeyboardInterrupt
#except Exception as e:
#    parser.exit(type(e).__name__ + ': ' + str(e))