#!/usr/bin/python3
# -*-coding:Utf-8 -*

''' From an stm reference file, generate a perfect ctm system file

    Usage: 
        
        python3 OpenASR_generate_ctm_file.py -f test_suite_validation/reference/doc1_inLine.stm -o tmp
'''

import argparse
import os
import pandas
import csv

def generateCTMFile(args):
	rows = []
	csv_header = ['Filename', 'Channel', 'Speaker','BeginTime','EndTime','Transcript']

	with open(args.STM_file, 'r') as f_input:
		for row in csv.DictReader(f_input, delimiter='\t', fieldnames=csv_header[:-1], restkey=csv_header[-1], skipinitialspace=True):
			try:
				rows.append([row['Filename'], row['Channel'], row['Speaker'], row['BeginTime'], row['EndTime'], ' '.join(row['Transcript'])])
			except Exception as e:
				rows.append([row['Filename'], row['Channel'], row['Speaker'], row['BeginTime'], row['EndTime'], ' '])

	df = pandas.DataFrame(rows, columns=csv_header)

	#Filter out the interSeg part
	df_subset = df.loc[~df['Speaker'].str.contains('interSeg')]

	pandas.set_option('mode.chained_assignment', None)
	#Get the duration of each word (endTime-beginTime/nb of words within transcript)
	df_subset['Dur']=(df_subset['EndTime'].astype(float)-df_subset['BeginTime'].astype(float))/df_subset['Transcript'].str.split(' ').str.len().astype(float)

	perfect_ctm = "{}.ctm".format(os.path.basename(args.STM_file).split(".")[0])

	if args.output_dir:
		output_file = os.path.join(args.output_dir, perfect_ctm)
		if os.path.exists(output_file):
			os.remove(output_file)
		else:			
			if os.path.exists(args.output_dir):
				pass
			else:
				os.mkdir(args.output_dir)
	else:
		current_directory = os.getcwd()
		output_file = os.path.join(current_directory, perfect_ctm)
		if os.path.exists(output_file):
			os.remove(output_file)

	#Iterate over each rows and each words with transcript (+ add Dur time after each iteration)
	with open(output_file,"w") as ctmFile:
		for index, row in df_subset.iterrows():
			for w in row['Transcript'].split(' '):
				ctmFile.write(row['Filename']+" "+row['Channel']+" "+str(row['BeginTime'])+" "+str(row['Dur'])+" "+w+"\n")
				row['BeginTime']=float(row['BeginTime'])+float(row['Dur'])
	
## Begin code
def main():
	parser = argparse.ArgumentParser(description='Generate a perfect CTM system file from an stm reference file')
	parser.add_argument('-f','--STM_file', required=True, help="STM reference file")
	parser.add_argument('-o','--output_dir', type=str, help='output directory containing the perfect corresponding CTM system file')
	parser.set_defaults(func=generateCTMFile)
	args = parser.parse_args()
	if hasattr(args, 'func') and args.func:
		args.func(args)
	else:
		parser.print_help()


if __name__ == '__main__':
	main()
