import pickle
import scipy.io.wavfile as wav
import wave
import numpy as np
from pprint import pprint
import sys
import lmdb 
import math
import os
from os import listdir
from os.path import isfile, join
from sklearn.preprocessing import normalize

caffe_root = '/home/erika/projects/caffe/'  

import sys
sys.path.insert(0, caffe_root + 'python')
import caffe
# Load data using parameters from a pre-created run file (generated by create_deepsound_net) 
# and save into a lightning database.

def prepare_data(run_name):
	
	with open(run_name + '.pickle') as f:  
		params = pickle.load(f)

	folder = params['folder']
	   
	data_dim = params['data_dim'] # how many samples to look at 
	test_instances = params['test_instances'] # number of tests for the test phase
	training_instances = params['training_instances'] # training phase instances from each sound file.
	fft = params['fft']
	raw2d = params['raw2d']
	raw2d_multiplier = params['raw2d_multiplier']


	all_data = np.zeros((1,  data_dim[0], data_dim[1]))
	all_labels = np.zeros((1))


	files = [f for f in listdir(folder) if isfile(join(folder, f))]
	for f in files:
		data = process_file(folder + '/' + f, data_dim, training_instances, fft, raw2d, raw2d_multiplier)
		if(len(data) > 0):
			code = int(f[:2])
			labels = np.empty((data.shape[0]))
			labels.fill(code)
			all_data = np.concatenate((all_data, data), axis=0)

			all_labels = np.concatenate((all_labels, labels), axis=0)


	all_data = all_data.reshape((all_data.shape[0], 1, data_dim[0], data_dim[1]))
	all_labels = all_labels.reshape((all_labels.shape[0], 1, 1, 1))

	print "all_data"
	print all_data.shape
	print "all_labels"
	print all_labels.shape


		

    # Shuffle data 
	rng_state = np.random.get_state()
	np.random.shuffle(all_data)
	np.random.set_state(rng_state)
	np.random.shuffle(all_labels)
	print "done shuffling."


	clip = np.empty(all_data.shape)
	clip.fill(1)
	#clip = np.insert(clip, 1, 0, axis=3)

	input_clip = clip[0:-test_instances]
	input_data = all_data[0:-test_instances]
	input_labels = all_labels[0:-test_instances]
	test_data = all_data[-test_instances:]
	test_labels = all_labels[-test_instances:]
	test_clip = clip[-test_instances:]
	print "input_data:"
	print input_data.shape
	print "input_labels:"
	print input_labels.shape
	print "test_data"
	print test_data.shape
	print "test_labels:"
	print test_labels.shape


	print "making databases..."


	make_database(test_labels, 'test_labels_lmdb', 1)
	make_database(test_data, 'test_inputs_lmdb', data_dim)
	make_database(input_data, 'inputs_lmdb', data_dim)
	make_database(input_labels, 'labels_lmdb', 1)
	#make_database(input_clip, 'clip_lmdb', 1)
	#make_database(test_clip, 'test_clip_lmdb', 1)

    
# process a file for use with the database
def process_file(filename, data_dim, training_instances, fft=True, raw2d=False, raw2d_multiplier=None):
	print "processing file: " + filename
	data = load_wav(filename)
	h = data_dim[0]
	w = data_dim[1]
	



	if(raw2d):
		training_instances = min(training_instances, len(data)/w)
		all_data = []
		for t in range(training_instances):
			all_data += [wav_to_raw2d(data[t*w:(t+1)*w], h, raw2d_multiplier)]
		all_data = np.array(all_data, dtype=np.float32)
		return all_data
	if(not fft):
		training_instances = min(training_instances, len(data)/w)
		all_data = np.zeros((training_instances, h, w))
		for t in range(training_instances):
			for hh in range(h):
				all_data[t][hh][:] = np.roll(data[t*w:(t+1)*w], hh)
		all_data = np.array(all_data, dtype=np.float32)
		print all_data.shape
		return all_data
	w = w/2
	if(training_instances*w*h > len(data)):
		training_instances = int(len(data)/(w*h))
	
	new_len = w*h*training_instances

	new_len = min(w*training_instances*h, len(data) - (len(data) % (training_instances*h)))

	data = data[:new_len]

	data.resize(len(data)/w, w)
	if(fft):


		fftdata = np.fft.fft(data, data.shape[1], axis=1)
		mag = np.absolute(fftdata)
		phase = np.angle(fftdata)
		data = np.concatenate((mag, phase), axis=1)


	f =  data.shape[0]*data.shape[1]/(h*w*2)


	data =data.reshape(( f, h, w*2))
	return data

# make a lightning database
def make_database(inputs, name, data_dim):
	raw_input("deleting and re-creating " + name + " [Y]")
	os.system('rm -r ' + name)
	in_db = lmdb.open(name, map_size=int(1e12))
	with in_db.begin(write=True) as in_txn:
		for in_idx, in_ in enumerate(inputs):
			im_dat = caffe.io.array_to_datum(in_)
			in_txn.put('{:0>10d}'.format(in_idx), im_dat.SerializeToString())
	in_db.close()




# returns magnitude and phase spectrum of the given data
def get_fft(data):
	fftdata = np.fft.fft(data)
	mag = np.absolute(fftdata)
	phase = np.angle(fftdata)
	return(mag, phase)


def mag_phase_to_complex(mag, phase):
	return np.array([mag[p]*np.exp(1.j*phase[p]) for p in range(0,len(mag))])


# def save_dream_wav(data, data_dim, filename):
# 	print "data"
# 	print len(data)
# 	print "data_dim:"
# 	print data_dim
# 	h = data_dim[0]/2
# 	w = data_dim[1]/2

# 	total_data = np.empty(())
# 	for i in range(len(data)):
	
# 		d = data[i][:]


# 		for q in range(0, d.shape[0]):
# 			dd = d[q][:]

# 			dd = mag_phase_to_complex(dd[:w], dd[w:]) 
# 			dd = np.fft.irfft(dd, w)

# 			total_data = np.append(total_data, dd)

# 			print total_data



# 	save_wav(total_data, filename)


def save_dream_wav(data, data_dim, filename, fft=True, raw2d=False, raw2d_multiplier=None):
	print "data"
	print len(data)
	h = data_dim[0]
	w = data_dim[1]/2
	total_data = np.empty(())
	for i in range(len(data)):
		print "i = " + str(i)
		d = data[i][:]
		print d.shape

		if(raw2d):
			dd = raw2d_to_wav(d, raw2d_multiplier)
			total_data = np.append(total_data, dd)
		else:
			for j in range(d.shape[0]):
				dd = d[j][:]
				#print dd.shape
				if(fft):
					dd = mag_phase_to_complex(dd[:w], dd[w:]) 
					dd = np.fft.irfft(dd, w)
				
					

				#print dd.shape

				#dd -= np.average(dd)
				#dd *= (2**15-1.)/np.amax(dd)
				total_data = np.append(total_data, dd)


	print "total_data:"
	print total_data.shape
	save_wav(total_data, filename)


# correct solution:
def softmax(x):
    """Compute softmax values for each sets of scores in x."""
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum(axis=0) # only difference

# takes a 2d plot of raw audio and converts it back to raw
def raw2d_to_wav(image, multiplier):

	
	rangeobj = np.empty(image.shape)
	for r in range(rangeobj.shape[0]):
		rangeobj[r][:]  = r
	
	raw = np.argmax(image, axis=0)

	#raw = np.sum(softmax(image)*rangeobj,axis=0)


	raw = (0.1/multiplier)*(raw - image.shape[0]/2) 
	print raw
	return raw


# takes raw data and returns a 2d plot 
def wav_to_raw2d(data, height, multiplier):
	data = data*multiplier
	data = data.astype(np.int16)
	data += height/2 
	data[data<0] = 0
	data[data>=height] = height - 1
	plot = np.zeros((height, data.size))
	for d in range(data.size):
		plot[data[d]][d] = 1
	return plot

def test_raw2d():
	data = load_wav('/home/erika/Music/songsinmyhead/17vacation (2016_11_04 09_58_27 UTC).wav')
	rawplot = wav_to_raw2d(data[500000:800000], 100, 0.001)
	returning = raw2d_to_wav(rawplot, 0.001)
	save_wav(returning, 'test_raw2d.wav')


# saves sound data to file
def save_wav(total_data,filename):


	total_data = total_data.astype(np.int16)
	z = np.zeros(total_data.shape, dtype=np.int16)
	total_data = np.append(total_data, z)

	rate = 44100
	output_file = wave.open(filename, "w")
	output_file.setparams((1, 2, rate, 0, "NONE", "not compressed"))
	output_file.writeframes(total_data)
	output_file.close()
	print "wrote to file: " + filename


# load wav file and convert to float32 mono. Does not normalize.
def load_wav(filename):
    data = wav.read(filename)
    data = np.array(data[1], dtype=np.float32)

    # Convert to mono if need be
    if(data.ndim > 1):
        data = np.sum(data, axis=1)/2.0

    data = data.reshape((data.size))

    return data



def test_sound_functions():
	data = load_wav('/home/erika/Music/songsinmyhead/17vacation (2016_11_04 09_58_27 UTC).wav')
	mag, phase = get_fft(data[0:200000])
	rdata = reverse_fft(mag, phase)
	print "rdata.shape:"
	print rdata.shape
	save_wav(rdata, 'test4.wav')



def test_sound_functions2():
	data = load_wav('/home/erika/Music/songsinmyhead/17vacation (2016_11_04 09_58_27 UTC).wav')
	print "getting fft:"
	fftdata = np.fft.fft(data[0:200000], 200000, axis=0)
	data = np.absolute(fftdata)
	print "reversing:"
	rdata = np.fft.irfft(data)

	print "rdata.shape:"
	print rdata.shape
	save_wav(rdata, 'test5.wav')


def test_save_dream_wav():


	data_dim = (600, 1000)
	inputs = process_file('/home/erika/Music/songsinmyhead/b/01blame (2016_11_12 15_12_36 UTC).wav', data_dim,10)
	save_dream_wav(inputs, data_dim, 'test_dream_wav.wav')



if __name__ == "__main__":
	print "testing raw 2d..."
	test_raw2d()