import streamlit as st
import torch
from tqdm import tqdm
from transformers import pipeline
import numpy as np
import time
import string


# JS
import nltk
nltk.download('wordnet')
from nltk.corpus import wordnet as wn
from nltk.tokenize import word_tokenize

@st.cache_resource
def get_models_context(llama=False):
  st.write('Loading the model...')
  model = pipeline("fill-mask")
  st.write("_The assistant is loaded and ready to use! :tada:_")
  return model 

model_context = get_models_context()

def remove_punctuation(word):
    # Create a translation table that maps all punctuation characters to None
    translator = str.maketrans('', '', string.punctuation)
    
    # Use the translate method to remove punctuation from the word
    word_without_punctuation = word.translate(translator)
    
    return word_without_punctuation

def return_top_k_context(sentence):

  if sentence[-1] != ".":
    sentence = sentence + "."
    
  output = model_context(sentence) 
  output = [output[i]['token_str'].strip() for i in range(len(output))]
  return output


def get_text():
  input_text = st.chat_input()
  return input_text

def write_bot(input, remember=True, blink=True):
    with st.chat_message('assistant'):
      message_placeholder = st.empty()
      full_response = input
      if blink == True:
        response = ''
        for chunk in full_response.split():
          response += chunk + " "
          time.sleep(0.05)
          # Add a blinking cursor to simulate typing
          message_placeholder.markdown(response + "▌") 
        time.sleep(0.5)
      message_placeholder.markdown(full_response)
    if remember == True:
      st.session_state.messages_context.append({'role': 'assistant', 'content': full_response})
    
def ask_if_helped_context():
  y = st.button('Yes!', key=60)
  n = st.button('No...', key=61)
  new = st.button('I have a new word', key=62)
  if y:
    write_bot("I am happy to help!")
    again = st.button('Play again')
    if again:
      write_bot("Please give a sentence using a <mask> instead of the word you have in mind!")
    st.session_state.is_helpful_context['ask'] = False
  elif n:
    st.session_state.actions_context.append('cue')
    st.session_state.is_helpful_context['ask'] = False
    #cue_generation()
  elif new:
    write_bot("Please give a sentence using a <mask> instead of the word you have in mind!")
    st.session_state.is_helpful_context['ask'] = False


def postproc_wn(related_words, syns=False):
    if syns:
        related_words = [word.split('.')[0] if word[0] != "." else word.split('.')[1] for word in related_words]
    else:
        related_words = [word.name().split('.')[0] if word.name()[0] != "." else word.name().split('.')[1] for word in related_words]
    related_words = [word.replace("_", " ") for word in related_words]

    return related_words
    
if 'messages_context' not in st.session_state:
  st.session_state.messages_context = []
  
if 'results_context' not in st.session_state:
  st.session_state.results_context = {'results_context': False, 'results_context_print': False}
  
if 'actions_context' not in st.session_state:
  st.session_state.actions_context = [""]
  
if 'counter_context' not in st.session_state:
  st.session_state.counter_context = {"letter_count": 0, "word_count": 0}
  
if 'is_helpful_context' not in st.session_state:
  st.session_state.is_helpful_context = {'ask':False}
  
if 'descriptions_context' not in st.session_state:
  st.session_state.descriptions_context = []
  
st.title("You name it! 🗣")

with st.chat_message('user', avatar='julian.jpg'):
  st.write("Hey assistant!")
  
bot = st.chat_message('assistant')
bot.write("Hello human! Wanna practice naming some words?")

#for showing history of messages_context
for message in st.session_state.messages_context:
  if message['role'] == 'user':
    with st.chat_message(message['role'], avatar='julian.jpg'):
      st.markdown(message['content'])
  else:
    with st.chat_message(message['role']):
      st.markdown(message['content'])
  
#display user message in chat message container
prompt_context = get_text() 
if prompt_context:
  with st.chat_message('user', avatar="julian.jpg"):
    st.markdown(prompt_context)
  #add to history
  st.session_state.messages_context.append({'role': 'user', 'content': prompt_context})
  #TODO: replace it with zero-shot classifier
  yes = ['yes', 'again', 'sure', 'new word', 'yes!', 'yep', 'yeah']
  try:
    if prompt_context.lower() in yes:
        write_bot("Please give a sentence using a <mask> instead of the word you have in mind!")
    #if previously we asked to give a prompt_context
    elif (st.session_state.messages_context[-2]['content'] == "Please give a sentence using a <mask> instead of the word you have in mind!") & (st.session_state.messages_context[-1]['content'] != "no"):
        write_bot("Great! Let me think what it could be...")
        st.session_state.descriptions_context.append(prompt_context)
        st.session_state.results_context['results_context'] = return_top_k_context(st.session_state.descriptions_context[-1])
        st.session_state.results_context['results_context_print'] = dict(zip(range(1, len(st.session_state.results_context['results_context'])+1), st.session_state.results_context['results_context']))
        write_bot("I think I have some ideas. Do you want to see my guesses or do you want a cue?")
        st.session_state.actions_context.append("result")
  except:
      write_bot("Sorry, I didn't understand you... I am still learning :sob: For now, could you respond with 'yes' or 'no'? ")
      
if st.session_state.actions_context[-1] == "result":
  col1, col2, col3, col4, col5 = st.columns(5)
  with col1:
    a1 = st.button('Results', key=10)
  with col2:
    a2 = st.button('Cue', key=11)
  if a1:
    write_bot("Here are my guesses about your word:")
    st.write(st.session_state.results_context['results_context_print'])
    time.sleep(1)
    write_bot('Does it help you remember the word?', remember=False)
    st.session_state.is_helpful_context['ask'] = True
  elif a2:
    write_bot(f'The first letter is {st.session_state.results_context["results_context"][0][0]}.')
    time.sleep(1)
    # st.session_state.actions_context.append('cue')
    #cue_generation()
    write_bot('Does it help you remember the word?', remember=False)
    st.session_state.is_helpful_context['ask'] = True
    
if st.session_state.is_helpful_context['ask'] == True:
    ask_if_helped_context()

if st.session_state.actions_context[-1] == 'cue':
  guessed = False
  write_bot('What do you want to see?', remember=False, blink=False)
  col1, col2, col3, col4, col5 = st.columns(5)


  with col1:
    b1 = st.button("Next letter", key="1")
  with col2:
    b2 = st.button("Related words")
  with col3:
    b3 = st.button("Next word", key="2")
  with col4:
    b4 = st.button("All words", key="3")

  b5 = st.button("I remembered the word!", key="4", type='primary')
  b6 = st.button("Exit", key="5", type='primary')
  new = st.button('Play again', key=64, type='primary')
  while guessed == False:
    # JS
    word_count = st.session_state.counter_context["word_count"]
    target = st.session_state.results_context["results_context"][word_count]

    if b1:
      st.session_state.counter_context["letter_count"] += 1
      letter_count = st.session_state.counter_context["letter_count"]
      if letter_count < len(target):
        write_bot(f'The word starts with {st.session_state.results_context["results_context"][word_count][:letter_count]}.', remember=False)
        st.session_state.is_helpful_context['ask'] = True
      else:
        write_bot(f'This is my predicted word: "{target}". Does this match your query?')
        st.session_state.is_helpful_context['ask'] = True

    elif b2:
      rels = return_top_k_context(st.session_state.descriptions_context[-1], word=target, rels=True)
      write_bot(f'Here are words that are related to your word: {", ".join(rels)}.', remember=False)
      st.session_state.is_helpful_context['ask'] = True

    elif b3:
      st.session_state.counter_context["letter_count"] = 1
      letter_count = st.session_state.counter_context["letter_count"]
      st.session_state.counter_context["word_count"] += 1
      word_count = st.session_state.counter_context["word_count"]
      if letter_count < len(target):
        write_bot(f'The next word starts with {st.session_state.results_context["results_context"][word_count][:letter_count]}.', remember=False)
        st.session_state.is_helpful_context['ask'] = True
      else:
        write_bot(f'This is my predicted word: "{target}". Does this match your query?')
        st.session_state.is_helpful_context['ask'] = True

    elif b4:
      write_bot(f"Here are all my guesses about your word: {st.session_state.results_context['results_context_print']}")
      st.session_state.is_helpful_context['ask'] = True
      
    elif b5:
      write_bot("Yay! I am happy I could be of help!")
      st.session_state.counter_context["word_count"] = 0
      st.session_state.counter_context["letter_count"] = 0
      new = st.button('Play again', key=63)
      if new:
        write_bot("Please give a sentence using a <mask> instead of the word you have in mind!")
      guessed = True
    
      break
    
    elif b6:
      write_bot("I am sorry I couldn't help you this time. See you soon!")
      st.session_state.counter_context["word_count"] = 0
      st.session_state.counter_context["letter_count"] = 0
    st.session_state.actions_context.append('cue')
    
    
    if new:
      write_bot("Please give a sentence using a <mask> instead of the word you have in mind!")
      st.session_state.counter_context["word_count"] = 0
      st.session_state.counter_context["letter_count"] = 0
    break