from scipy.io import wavfile
import numpy as np
import os, sys, time
sys.path.append(os.path.abspath(os.path.join(os.getcwd(), "../../")))
from args import args
from model import params, wavenet
import data

def generate_audio():
	# compute receptive field width
	learnable_steps = 1
	batch_size = 1
	num_layers = len(params.residual_conv_channels)
	receptive_steps_per_unit = params.residual_conv_filter_width ** num_layers
	receptive_steps = (receptive_steps_per_unit - 1) * params.residual_num_blocks + 1
	target_width = learnable_steps
	input_width = receptive_steps
	# to compute all learnable targets
	input_width += learnable_steps - 1
	## padding for causal conv block
	input_width += len(params.causal_conv_channels)

	# quantized signals generated by WaveNet
	generated_quantized_audio = np.zeros((input_width, ), dtype=np.int32)

	start = time.time()
	for time_step in xrange(9):
		# quantized signals in receptive field
		padded_quantized_x_batch = generated_quantized_audio[-input_width:].reshape((1, -1))

		# convert to image
		padded_x_batch = data.onehot_pixel_image(padded_quantized_x_batch, quantization_steps=params.quantization_steps)

		# generate next signal
		softmax = wavenet.forward_one_step(padded_x_batch, softmax=True, return_numpy=True)
		softmax = softmax[0, :, 0, -1]
		generated_quantized_signal = np.argmax(softmax)
		generated_quantized_audio = np.append(generated_quantized_audio, [generated_quantized_signal], axis=0)
		print generated_quantized_signal,

	print generated_quantized_audio
	print time.time() - start

	wavenet.prev_causal_outputs = None
	wavenet.prev_residual_outputs_out = None
	wavenet.prev_residual_outputs_z = None
	generated_quantized_audio = np.zeros((input_width, ), dtype=np.int32)

	start = time.time()
	for time_step in xrange(9):
		# quantized signals in receptive field
		padded_quantized_x_batch = generated_quantized_audio[-input_width:].reshape((1, -1))

		# convert to image
		padded_x_batch = data.onehot_pixel_image(padded_quantized_x_batch, quantization_steps=params.quantization_steps)

		# generate next signal
		softmax = wavenet._forward_one_step(padded_x_batch, softmax=True, return_numpy=True)
		softmax = softmax[0, :, 0, -1]
		generated_quantized_signal = np.argmax(softmax)
		generated_quantized_audio = np.append(generated_quantized_audio, [generated_quantized_signal], axis=0)
		print generated_quantized_signal,

	print generated_quantized_audio
	print time.time() - start

def main():
	generate_audio()

if __name__ == '__main__':
	main()
