from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModel
import json
import spacy

from celebbot import CelebBot
from utils import *

DEBUG = False
QA_MODEL_ID = "google/flan-t5-large"
SENTTR_MODEL_ID = "sentence-transformers/all-mpnet-base-v2"

def main():
    with open("data.json") as json_file:
        celeb_data = json.load(json_file)
    message = "Please choose your favorite celebrity from\n"\
              "1. Cate Blanchett\n"\
              "2. David Beckham\n"\
              "3. Emma Watson\n"\
              "4. Lady Gaga\n"\
              "5. Madonna\n"\
              "6. Mark Zuckerberg\n"\
              "input name:\n"
    QA_tokenizer = AutoTokenizer.from_pretrained(QA_MODEL_ID)
    QA_model = AutoModelForSeq2SeqLM.from_pretrained(QA_MODEL_ID)
    sentTr_tokenizer = AutoTokenizer.from_pretrained(SENTTR_MODEL_ID)
    sentTr_model = AutoModel.from_pretrained(SENTTR_MODEL_ID)

    name = input(message)
    gender = celeb_data[name]["gender"]
    knowledge = celeb_data[name]["knowledge"]
    
    lname = name.split(" ")[-1]
    lname_regex = re.compile(rf'\b({lname})\b')
    name_regex = re.compile(rf'\b({name})\b')
    lnames = lname+"’s" if not lname.endswith("s") else lname+"’"
    lnames_regex = re.compile(rf'\b({lnames})\b')
    names = name+"’s" if not name.endswith("s") else name+"’"
    names_regex = re.compile(rf'\b({names})\b')
    if gender == "M":
        knowledge = re.sub(he_regex, "I", knowledge)
        knowledge = re.sub(his_regex, "my", knowledge)
    elif gender == "F":
        knowledge = re.sub(she_regex, "I", knowledge)
        knowledge = re.sub(her_regex, "my", knowledge)
    knowledge = re.sub(names_regex, "my", knowledge)
    knowledge = re.sub(lnames_regex, "my", knowledge)
    knowledge = re.sub(name_regex, "I", knowledge)
    knowledge = re.sub(lname_regex, "I", knowledge)
    
    spacy_model = spacy.load("en_core_web_sm")
    knowledge_sents = [i.text.strip() for i in spacy_model(knowledge).sents]

    ai = CelebBot(name, QA_tokenizer, QA_model, sentTr_tokenizer, sentTr_model, spacy_model, knowledge_sents)

    answers = []
    
    while True:
    # for q in celeb_data[ai.name_id][ai.name]["questions"]:
        if not DEBUG:
            ai.speech_to_text()
        else:
            ai.text = input("Your question: ")
            pass

        # ai.text = q
        if ai.text != "":
            print("me --> ", ai.text)

            answers.append(ai.question_answer())

            if not DEBUG:
                ai.text_to_speech()
        ai.text = ""

# Run the AI
if __name__ == "__main__":
    main()