from modules.NLU import NLU
from modules.DST import DST
#from modules.DP import DP
#from modules.NLG import NLG
import re

def format_prediction(prediction, intent):
    out_list = []
    for idx, tup in enumerate(prediction):
        if tup[1][0] == 'B':
            slot_list = [intent, 'Cinema', tup[1][2:], tup[0]]
            for tup in prediction[idx + 1:]:
                if tup[1][0] != 'I':
                    break
                else:
                    slot_list[3] += ' ' + tup[0]
            out_list.append(slot_list)
    for slot in out_list:
        slot[3] = re.sub("^[!\"#$%&\'()*+,.;:<=>?\[\]^_`{|}~]+", '', slot[3])
        slot[3] = re.sub("[!\"#$%&\'()*+,.;:<=>?\[\]^_`{|}~]+$", '', slot[3])
    return out_list

def main():
    nlu = NLU()
    dst = DST()
    #dp = DP()
    #nlg = NLG()
    
    nlu.train_slot_model('data/train+test-pl.conllu', 'data/train+test-pl.conllu')
    nlu.train_intent_model('data/NLU_data_intent')
    # nlu.load_slot_model('slot-model-pl')
    # nlu.load_intent_model('intent-model-pl')

    print('===========================================')
    print('### By otrzymać pomoc, wpisz /pomoc ###')
    print('### By zakończyć rozmowę, wpisz /koniec ###')
    print('Witaj, jestem Usher - system do rezerwacji biletów kinowych. W czym mogę Ci pomóc?')
        
    # WIP
    while True:
        user_input = input('> ')
        if user_input == '/pomoc':
            print('TEKST_POMOCY_WIP')
        elif user_input == '/koniec':
            print('Dziękuję za skorzystanie z moich usług. Miłego dnia!')
            break
        else:
            slots = nlu.predict_slots(user_input)
            intent = nlu.predict_intent(user_input)
            formatted_prediction = format_prediction(slots, intent)
            print(formatted_prediction) # NLU output
            dst.update(formatted_prediction)
            print(dst.state) # DST output
            #DP, NLG...
            
    
if __name__ == '__main__':
    main()