from wordcloud import WordCloud, STOPWORDS, ImageColorGenerator
import matplotlib.pyplot as plt
from PIL import Image
import streamlit as st

def display_wordcloud(answer, answer_str):
    wc_all, wc_question, wc_reference  = st.columns([3, 3, 3])
    wordcloud = WordCloud(max_font_size=50, max_words=1000, background_color="white")
    with wc_all:
        image = Image.open('docs/images/all_papers_wordcloud.png')
        st.image(image)
        st.caption('''###### Corpus term frequecy.''')
    with wc_question:
        wordcloud_q = wordcloud.generate(answer_str)
        st.image(wordcloud_q.to_array())
        st.caption('''###### Answer term frequecy.''')
    with wc_reference:
        all_reference_texts = ''
        for nodewithscore in answer.source_nodes:  
            node = nodewithscore.node
            from llama_index.core.schema import NodeRelationship
                    #if NodeRelationship.SOURCE in node.relationships:
            all_reference_texts = all_reference_texts + '\n' + node.text
        wordcloud_r = wordcloud.generate(all_reference_texts)
        st.image(wordcloud_r.to_array())
        st.caption('''###### Reference plus graph term frequecy.''')