import argparse, json, sys, time, random
import spacy
from aitextgen import aitextgen
import string

def clean(text: str) -> str:

	s = text.split('\n')

	if(len(s) > 0):
		tok_1 = s[0].split(' ')
		if len(tok_1) > 0 and tok_1[0].strip() in string.punctuation:
				s_1 = ' '.join(tok_1[1:])
				s[0] = s_1.capitalize()	
		else:
			s[0] = s[0].capitalize()

	return '\n'.join(s)

def format(text: str) -> str:

	return text.replace('\r\n', '\n').replace('\n\n', '\n')
				
def main() -> int:

	p = argparse.ArgumentParser()
	p.add_argument("-c", "--config", type=str, default="config.json", help="configuratin file")
	p.add_argument("-i", "--iterations", type=int, default=10, help="number of iterations")
	args = p.parse_args()

	print(args)

	with open(args.config) as f:
		conf = json.load(f)

	voices = []
	for v in conf['voices']:
		a = aitextgen(model_folder=v['model_dir'], tokenizer_file=v['tokeniser_file'])
		voices.append({"name": v["name"].upper(), "a": a, "temp": float(v["temperature"])})

	nbr_voices = len(voices)
	current_voice = ""
	for i in range(args.iterations):
		rindex = random.randint(0, nbr_voices - 1)
		v = voices[rindex]
		if v['name'] != current_voice:
			print("==========")
			print(v['name'] + ":")
			current_voice = v['name']
		t = v['a'].generate(n=1, max_lenght=32, temperature=v['temp'], return_as_list=True)[0]
		if t != None:
			t = clean(t)
			t = format(t)
			print(t)
		
		time.sleep(4)


if __name__ == '__main__':
    sys.exit(main())