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

		python OPENASR21_convert_reference_transcript.py \
		-l Cantonese
		-i BABEL_BP_101_10470_20111118_172644_inLine.txt \
		-c 1
		-o tmp
'''

import os, re
import argparse
import pandas as pd
import csv

def tokenize(label, language):

	#remove some tags, marks, characters and punctuations if any
	language = language.lower()
	if language == "farsi":
		label_punc_rm = re.sub('\.(?!\d)|\,|\?|\u200c|\=|\"|\!', '', label)
	else:
		label_punc_rm = re.sub('\.(?!\d)|\,|\?|\u200c|\"|\!', '', label)

	characters_rm = ["(())",
				"<no-speech>",
				"--",
				"~",
				"<sta>",
				"%incomplete",
				"<lipsmack>", "<breath>", "<cough>", "<laugh>", "<click>", "<ring>", "<dtmf>", "<int>", "<misc>", "<dtmf>",
				"<male-to-female>", "<female-to-male>"]

	label_char_rm = label_punc_rm
	for character in characters_rm:
		label_char_rm = label_char_rm.replace(character, "")

	words = label_char_rm.split()

	#convert optional deletable tag
	def optional_deletable_space(word):

		if word == "<hes>" or word == "<foreign>" or re.match('^\-.*', word) or re.match('.*\-$', word):
			word = "(%s)" % word
		if re.match('\*.*\*',word):
			word = "(%s)" % word.strip('*')
		if "_" in word and word != "_":
			word = re.sub("[\_]"," ",word)
		if re.match('\/.*\/', word):
			word = "%s" % word.strip('/')

		return word

	words_tokenized = list(map(optional_deletable_space, words))

	if len(words_tokenized) > 0:
		#replace with the string "IGNORE_TIME_SEGMENT_IN_SCORING" when <overlap> or <prompt> appears.
		if "<overlap>" in words_tokenized or "<prompt>" in words_tokenized:
			label_tokenized = "IGNORE_TIME_SEGMENT_IN_SCORING"
		else:
			label_tokenized = " ".join(words_tokenized)
	else:
		label_tokenized = ""

	return label_tokenized

def txt_to_stm(df, filename, channel_name, language):

	time = []
	label = []

	i, j = 0, 1

	while (i < len(df.content) + 1) and (j < len(df.content) + 1):
		try:
			time.append(df.content[i])
			label.append(df.content[j])
		except:
			pass
		i += 2
		j += 2

	time = [re.sub(r"[\[\]]", "", a) for a in time]

	dict_new = {"time" : time, "label" : label}
	df_new = pd.DataFrame({key:pd.Series(value) for key, value in dict_new.items()})

	records = []

	for v, w in zip(df_new["time"][:-1], df_new["time"][1:]):

		if isinstance(df_new[df_new['time'] == v]['label'].item(),str):
			tokenized_text = tokenize(df_new[df_new['time'] == v]['label'].item(), language)
		else:
			tokenized_text = tokenize(str(df_new[df_new['time'] == v]['label'].item()), language)

		#if normalized transcript is empty, speakerid will be added interSeg suffix
		if tokenized_text == "":
			speaker_name_new = "{}_{}interSeg".format(filename, channel_name)
			record = [filename, channel_name, speaker_name_new, v, w]

		else:
			record = [filename, channel_name, "{}_{}".format(filename, channel_name), v, w, tokenized_text]

		record = ['{:0.3f}'.format(x) if isinstance(x, float) else x for x in record]
		records.append(record)

	return records

def convert_from_transcript(args):

	file_base = os.path.basename(args.input_transcript).split(".")[0]

	if file_base.split("_")[-1] == "inLine" or file_base.split("_")[-1] == "outLine":
		new_file_base = "_".join(file_base.split("_")[:-1])
	else:
		new_file_base = file_base

	print("Converting {} transcript into stm format".format(args.input_transcript))
	transcript_df = pd.read_csv(args.input_transcript, sep = "\n", header = None, names = ["content"], quoting = 3)
	result = txt_to_stm(transcript_df, new_file_base, str(args.channel), args.language)

	transcript_stm = "{}.stm".format(file_base)

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

	with open(output_file, "w") as output:
		writer = csv.writer(output, delimiter = '	')
		writer.writerows(result)
	output.close()

def main():
	parser = argparse.ArgumentParser(description='Convert OPENASR21 transcript into stm format')
	parser.add_argument('-l','--language', type=str, required=True, help='indicates the language of the input transcript file (e.g., Mandarin)')
	parser.add_argument('-i','--input_transcript', type=str, required=True, help='input file containing the transcript in dev/train data')
	parser.add_argument('-c','--channel', type=str, required=True, help='indicates the channel ID. If the file is part of conversational telephone speech, the inline channel is designated as 1 while the outline channel is designated as 2. Otherwise, the channel ID is 1.')
	parser.add_argument('-o','--output_directory', type=str, help='output directory containing the transcript in STM format')

	parser.set_defaults(func=convert_from_transcript)

	args = parser.parse_args()

	if hasattr(args, 'func') and args.func:
		args.func(args)
	else:
		parser.print_help()

if __name__ == '__main__':
	main()
