from aitextgen import aitextgen
import utterance.utils
import gensim, regex, string, time

UTTERANCE_MEMORY_LEN = 15
UTTERANCE_MEMORY_MIN_DIST = 0.85

class Voice:

	def __init__(self, name: str, model: str, tokenizer: str, temp: int, length: int):
		self.name = name
		self.temp = temp
		self.length = length
		self.v = aitextgen(model_folder=model, tokenizer_file=tokenizer)
		self.utterances = []

	def utter_n(self, n: int, temp: float = None, length: int = None):
		t = self.temp if temp == None else temp
		l = self.length if length == None else length
		return self.v.generate(n=n, max_length=l, temperature=t, seed=int(time.time()), return_as_list=True)

	def utter_one(self, temp: float = None, length: int = None) -> str:
		t = self.temp if temp == None else temp
		l = self.length if length == None else length		
		return self.utter_n(n=1, temp=temp, length=length)[0]

	def prompt(self, pinput: str, temp: float = None, length: int = None) -> str:
		t = self.temp if temp == None else temp
		l = self.length if length == None else length
		return self.v.generate(prompt=pinput, n=1, max_length=l, temperature=t, seed=int(time.time()), return_as_list=True)[0]

	def select(self, utterance: str) -> bool:

		# fuction making sure to not say the same thing

		toks = set(gensim.utils.tokenize(utterance))
		
		i = 0
		for u in self.utterances:
			d = gensim.matutils.jaccard_distance(toks, u)
			print(f"{i} - d: {d}")
			if d < UTTERANCE_MEMORY_MIN_DIST:
				return False

		self.utterances.append(toks)

		if len(self.utterances) > UTTERANCE_MEMORY_LEN:
			self.utterances.pop(0)

		return True

	def remember(self, utterance: str):
		toks = set(gensim.utils.tokenize(utterance))
		self.utterances.append(toks)


	def set_channel(self, root: str, endpoint: str):
		self.channel = root + endpoint

	def channel(self):
		return self.channel

	def set_calculate(self, root: str, endpoint: str):
		self.calculate = root + endpoint

	def calculate(self):
		return self.calculate

	def set_temperature(self, root: str, endpoint: str):
		self.temperature = root + endpoint

	def calculate(self):
		return self.temperature

	def fragments(self, utt: str):
		self.utterance = utt
		self.utterance_fragments = utterance.utils.fragments(utt)
		return self.utterance_fragments