from scipy.io import wavfile
import numpy as np
import os, sys, time
sys.path.append(os.path.split(os.getcwd())[0])
from args import args
from model import params, wavenet
import data

def generate_audio(sampling_rate=48000, generate_sec=1, remove_silence_frames=False):
	batch_size = 1

	# compute required input width
	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
	input_width = receptive_steps
	# 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.time()
	for time_step in xrange(1, int(sampling_rate * generate_sec)):
		# 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
		if args.use_faster_wavenet:
			softmax = wavenet._forward_one_step(padded_x_batch, softmax=True, return_numpy=True)
		else:
			softmax = wavenet.forward_one_step(padded_x_batch, softmax=True, return_numpy=True)

		softmax = softmax[0, :, 0, -1]
		generated_quantized_signal = np.random.choice(np.arange(params.quantization_steps), p=softmax)

		if generated_quantized_signal == 0 and remove_silence_frames:
			pass
		else:
			generated_quantized_audio = np.append(generated_quantized_audio, [generated_quantized_signal], axis=0)

		if time_step % 10 == 0:
			sys.stdout.write("\rgenerating {:.2f} msec / {:.2f} msec".format(time_step * 1000.0 / sampling_rate, generate_sec * 1000.0))
			sys.stdout.flush()

	print "\ndone in {:.3f} sec".format(time.time() - start_time)

	# remove zero paddings
	generated_quantized_audio = generated_quantized_audio[input_width:]

	try:
		os.mkdir(args.generate_dir)
	except:
		pass

	filename = "{}/generated.wav".format(args.generate_dir)
	data.save_audio_file(filename, generated_quantized_audio, params.quantization_steps, format="16bit_pcm", sampling_rate=sampling_rate)

def main():
	generate_audio(generate_sec=args.generate_sec, sampling_rate=params.sampling_rate)

if __name__ == '__main__':
	main()
