'''
  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('
  ').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 + "
  ")
  			else:
  				f.write(metas_str + "
  ")
  
  '''
  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")