'''
The goal of this script is to apply a clustering to pvector in order to find new classes assigned for each utterance or frame.
This new class can be used for new training systems, replacing character classes for example by the calculatering classes from clustering. 
We hope this will generate interesting classes that will help the system to understand the structure of the voices.

TODO: Change it in such a way as to take a number (1, 2, 3, 4) and calculate everything needed like clustering. Train on the train set and then project the test on this clustering in order to know to what cluster it belong to.
'''

import os
import numpy as np
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
import argparse
import pandas as pd
import pickle


'''
Return data in panda format version
'''
def read_vecfile(filepath, toy_version=False):
	vectors = ""
	metas = ""
	with open(filepath, "r") as f:
		for i, line in enumerate(f):
			if toy_version == True and i > 100:
				break
			spl_line = line.split(" ")
			
			if(len(vectors) == 0):
				vectors = np.empty((0, len(spl_line[1:])), np.float32)
				metas = np.empty((0, len(spl_line[0].split(","))))
			
			# Then we add the current line to the data
			metas = np.append(
				metas, 
				np.asarray([spl_line[0].split(",")]), 
				axis=0)
		
			vectors = np.append(
				vectors, 
				np.asarray([spl_line[1:]], dtype=np.float32), 
				axis=0)
	return (metas, vectors)

'''
Return list of metas of the listfile
'''
def read_lstfile(filepath, toy_version=False):
	metas = np.empty((0, 4))
	with open(filepath, "r") as f:
		for i, line in enumerate(f):
			if toy_version == True and i > 100:
				break
			metas = np.append(
				metas,
				np.asarray([line.rstrip('\n').split(",")]),
				axis=0)
	return metas
	
'''
Save a vector file from metas and vector values
'''
def save_file(filepath, metas, values=None):
	with open(filepath, "w") as f:
		for i  in range(len(metas)):
			metas_str = ",".join(str(v) for v in metas[i])
			if not values == None:
				try:
					infos_str = " ".join(str(v) for v in values[i])
				except TypeError as te:
					infos_str = str(values[i])
				f.write(metas_str + " " + infos_str + "\n")
			else:
				f.write(metas_str + "\n")

'''
Take the data and index them. 
'''
def index_data(metas, vectors):
	data = {}
	data["en-us"] = {}
	data["fr-fr"] = {}
	for i, vector in enumerate(vectors):
		meta = metas[i]
		data[meta[0]][meta[3]] = {}
		data[meta[0]][meta[3]]["metas"] = meta
		data[meta[0]][meta[3]]["vector"] = vector
	return data	

		

'''
Récupère un sous ensemble des données de base à partir d'une 
liste.
'''
def get_subdata(data, lst):
	metas = ""
	vectors = ""
	for meta in lst:
		vector = data[meta[0]][meta[3]]["vector"]
		if(len(metas) == 0):	
			metas = np.empty((0, len(meta)))
			vectors = np.empty((0, len(vector)), np.float64)
		metas = np.append(
			metas,
			np.asarray([data[meta[0]][meta[3]]["metas"]]),
			axis=0)
		vectors = np.append(
			vectors,
			np.asarray([vector]),
			axis=0)
	return metas, vectors	
	

'''
Apply clustering on data of filename. 
Use a list to determine train et test using train valid, test. 
Save the file with the given suffix. 
Check the existence of the files before calculating and saving,
if the two files already exist, it will not calculate it again. 

However: if one file is not present, this function will calculate 
it again. 

TODO: Add a variable to force the calculation of all the files 
even if they exist. 
'''
def apply_clustering(filename, dir_lst, dir_data, suffix_outfile):

	# Applicate it for normal version 	
	metas, vectors = read_vecfile(os.path.join(dir_data, filename), toy_version=False)
	data = index_data(metas, vectors)

	### CURSOR
	# Get Train 
	train_lst = read_lstfile(os.path.join(dir_lst, "train_" + str(NUMBER) + ".lst"))
	train_metas, train_vectors = get_subdata(data, train_lst)

	# Get Val 
	val_lst = read_lstfile(os.path.join(dir_lst, "val_" + str(NUMBER) + ".lst"))
	val_metas, val_vectors = get_subdata(data, val_lst)

	# Get Test
	test_lst = read_lstfile(os.path.join(dir_lst, "test_" + str(NUMBER) + ".lst"))
	test_metas, test_vectors = get_subdata(data, test_lst)
	
	# Verif shapes
	print("verif shapes")
	print(train_metas.shape)
	print(val_metas.shape)
	print(test_metas.shape)

	# Entrainer le k-means sur le train + val
	#Ks = [12, 24, 48]

	print("k=[", end="")
	Ks = [6,12,24,48,64]
	for k in Ks:
		# Process the name
		suffix = "_" + suffix_outfile if not suffix_outfile == "" else ""
		k_str = "{:03d}".format(k) # K in string
		filename_pickle = os.path.join(
			DIR_DATA, 
			"clusters_trained_on_train_" +str(k_str)+ "_pickle_" + suffix + ".txt")	
		filename_clusters = os.path.join(
			DIR_DATA, 
			"clusters_trained_on_train_" +str(k_str)+ suffix + ".txt")
		
		# Check if on of the two file does not exist
		condition = not(
			os.path.exists(filename_pickle) 
			and os.path.exists(filename_clusters)
		)
			
		if condition:	
			print(str(k)+",", end=" ")
			kmeans = KMeans(n_clusters=k, n_init=10, random_state=0).fit(
				train_vectors)	
			test_pred = kmeans.predict(np.concatenate((val_vectors, test_vectors), axis=0))
			metas_tosave = np.concatenate([train_metas, val_metas, test_metas], axis=0)
			values_tosave = np.concatenate([kmeans.labels_, test_pred], axis=0)
			metas_tosave[:, 1] = values_tosave # Replace char by clusters
			save_file(filename_clusters, metas_tosave)
			pickle.dump(kmeans, open( filename_pickle, "wb" ) )
	print("]")

for NUMBER in range(1, 5):
	print("JACKKNIFING NUMBER: " + str(NUMBER))
	DIR_MAIN="exp/pvector-1"
	DIR_DATA=os.path.join(DIR_MAIN, str(NUMBER))
	DIR_LST=os.path.join(DIR_MAIN, "lst")
	OUTFILE_NAME="clustering"

	print("Calculating mass_effect_pvectors")
	apply_clustering("masseffect_pvectors.txt", 
		dir_lst = os.path.join(DIR_MAIN, "lst"),
		dir_data = DIR_DATA,
		suffix_outfile = "")

	print("Calculating mass_effect_pvectors_final")
	apply_clustering("masseffect_pvectors_final.txt",
		dir_lst = os.path.join(DIR_MAIN, "lst"),
		dir_data = DIR_DATA,
		suffix_outfile = "final")