from wordcloud import WordCloud
import matplotlib.pyplot as plt
import json
import requests
import gradio as gr
from io import BytesIO
import base64
import os

AUTH = os.environ["AUTH"]

def generate_wordcloud(keywords, theme):
    # Sort keywords in descending order based on their values
    sorted_keywords = [entry['keyword'] for entry in sorted(keywords, key=lambda x: x['value'], reverse=True)]

    keyword_dict = {entry['keyword']: entry['value'] for entry in keywords}
    
    colour_range =['#00B3CF', '#FF0079', '#8C72B4', '#99B561', '#C49C8E']
    
    # Create a color function mapping each word to a color from the range based on grouped index
    def color_func(word, font_size, position, orientation, random_state=42, **kwargs):
        index = sorted_keywords.index(word) // 2 % len(colour_range)
        return colour_range[index]
    
    # Generate the word cloud
    wordcloud = WordCloud(width=800, height=400,
                          background_color='white',
                          color_func=color_func,
                          collocations=False,  # Avoid combining words
                          random_state=42
                         ).generate_from_frequencies(keyword_dict)
    
    # Save the word cloud image to a BytesIO object
    image_stream = BytesIO()
    wordcloud.to_image().save(image_stream, format='PNG')
    image_stream.seek(0)

    # Convert the BytesIO object to a base64-encoded image string
    image_data = base64.b64encode(image_stream.read()).decode('utf-8')
    image_src = f"data:image/png;base64,{image_data}"

    return f'<img src="{image_src}" alt="{theme} Wordcloud" style="width:270px;height:150px;margin-right:30px;">'


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    with gr.Row():
        answers = gr.TextArea(label="List of responses", placeholder='["this is cool", "this is very cool"]')
        with gr.Column():
            nr_topics = gr.Slider(minimum=5, maximum=15, step=1)
    greet_btn = gr.Button("Generate themes")
   
    output = gr.HTML(label="Themes and Wordclouds")
    

    @greet_btn.click(inputs=[answers, nr_topics], outputs=output)
    def greet(answers, nr_topics):
        inputs = {"answers": json.loads(answers), "nr_topics": nr_topics}
        p = requests.post(
            url='https://pulsifi-dev--thematic-analysis-app.modal.run/analysis',
            headers={"Authorization": AUTH},
            json=inputs
        )

        result = p.json()
        #print(result)
        data = result.get("data")
        if data["thematic_results"]:
            output_text = ""
            for theme_data in data["thematic_results"]:
                theme = theme_data["theme"]
                description = theme_data["description"]
                keywords_list = theme_data['keywords']

                # Generate word cloud and get HTML content
                wordcloud_html = generate_wordcloud(keywords_list, theme)

                # Add theme description to the output text
                output_text += f"""
                    <div style="display:flex; align-items: center;">
                        {wordcloud_html}
                        <div>
                            <h3 style="color:blue">{theme}:</h3>
                            <p style="text-align: justify;">{description.strip()}</p>
                            <br> <!-- Add a line break after the theme description -->
                        </div>
                    </div>
                """

        else:
            output_text = data["reason"]

        return output_text

if __name__ == "__main__":
    demo.launch()