#------------------------------------------------------------------------
# Import Modules
#------------------------------------------------------------------------

import streamlit as st
import string
from annotated_text import annotated_text
from PIL import Image
from pathlib import Path
from dotenv import load_dotenv
load_dotenv()

#------------------------------------------------------------------------
# Configurations
#------------------------------------------------------------------------

# Streamlit page setup
# icon = Image.open("MTSS.ai_Icon.png")
icon = Image.open("/Users/cheynelevesseur/Desktop/Python_Code/LLM_Projects/LLM_Prxmpting/MTSS.ai_Icon.png")
st.set_page_config(
    page_title="Kaleidoscope | Text Annotation", 
    page_icon=icon,
    layout="centered", 
    initial_sidebar_state="auto",
    menu_items={
        'About': "### *This application was created by*  \n### LeVesseur Ph.D | MTSS.ai"
    }
)

#------------------------------------------------------------------------
# Header
#------------------------------------------------------------------------

# st.image('MTSS.ai_Logo.png', width=300)

st.title('MTSS:grey[.ai]')
st.header('Kaleidoscope:grey[ | Text Annotation]')

contact = st.sidebar.toggle('Handmade by  \n**LeVesseur** :grey[ PhD]  \n| :grey[MTSS.ai]')
if contact:
    st.sidebar.write('Inquiries: [info@mtss.ai](mailto:info@mtss.ai)  \nProfile: [levesseur.com](http://levesseur.com)  \nCheck out: [InkQA | Dynamic PDFs](http://www.inkqa.com)') 

#------------------------------------------------------------------------
# Sidebar
#------------------------------------------------------------------------    
    
# Color options
colors = {
    "Green (DAF1E7)": "#DAF1E7",
    "Blue (BDE5FF)": "#BDE5FF",
    "Navy (D1DBE9)": "#D1DBE9",
    "Teal (D6EAED)": "#D6EAED",
    "Iceburg (E4EEF6)": "#E4EEF6",
    "Vermillion (F6DCDD)": "#F6DCDD",
}

with st.sidebar:
    st.divider()
    # Sidebar display (Option 1: Color blocks with hex)
    st.sidebar.header("Recommended Colors")

    for color_name, hex_code in colors.items():
        st.sidebar.color_picker(color_name, hex_code)
    
    st.subheader("Example")
        
    annotated_text(
        "I",
        " ",
        "really",
        " ",
        ("appreciate", "Emotion", "#daf1e7"),
        " ",
        "all",
        " ",
        "that",
        " ",
        "the",
        " ",
        "social",
        " ",
        "committee",
        " ",
        "has",
        " ",
        "done",
        " ",
        "to",
        " ",
        "keep",
        " ",
        "us",
        " ",
        "feeling",
        " ",
        "connected",
        " ",
        ".",
        " ",
        "I",
        " ",
        "also",
        " ",
        "really",
        " ",
        ("value", "Priority", "#bde5ff"),
        " ",
        "our",
        " ",
        "in",
        " ",
        "-person",
        " ",
        "meetings",
        " ",
        "and",
        " ",
        "the",
        " ",
        "social",
        " ",
        "opportunities",
        " ",
        "built",
        " ",
        "into",
        " ",
        "these",
        " ",
        "meetings",
        " ",
        ".",
    )
    
    st.divider()
   
    st.subheader("Directions for Using the Text Annotation Tool")

    directions = """
    1. **Enter Your Text**:
    - Type or paste the text you want to annotate in the text area provided.
    
    2. **Submit Your Text**:
    - Click the "Submit Text" button to process your input. This will split your text into individual words.

    3. **Annotate the Words**:
    - For each word, you can provide a label (e.g., "Emotion", "Priority", "Verb").
    - Choose a color for each label using the color picker.

    4. **Generate Annotated Text**:
    - Once you have labeled and colored the words you want to annotate, click the "Generate Annotated Text" button.
    - The annotated text will be displayed, showing your labels and chosen colors.

    5. **Take a Screenshot**:
    - To use the annotated text, take a screenshot of the displayed text.

    6. **Adjust Text Width** (Optional):
    - If you want to adjust the width of the sentences for a better screenshot, minimize or resize your browser window accordingly before taking the screenshot.
    """

    st.markdown(directions)

#------------------------------------------------------------------------
# Functions: Personalized Labels and Colors
#------------------------------------------------------------------------

# Function to split text into words
def split_text(text):
    # Add a space before punctuation marks
    for char in string.punctuation:
        text = text.replace(char, f" {char}")
    return text.split()

# Initialize session state to store text and annotations
if 'user_text' not in st.session_state:
    st.session_state.user_text = ""
if 'words' not in st.session_state:
    st.session_state.words = []
if 'labels' not in st.session_state:
    st.session_state.labels = []
if 'colors' not in st.session_state:
    st.session_state.colors = []

# User input for the text
user_text = st.text_area("Enter the text you want to annotate:", value=st.session_state.user_text, height=100)

# Button to process the text
if st.button("Submit Text"):
    st.session_state.user_text = user_text
    st.session_state.words = split_text(user_text)
    st.session_state.labels = [""] * len(st.session_state.words)
    st.session_state.colors = ["#FFFFFF"] * len(st.session_state.words)

# Collect annotation inputs for each word
if st.session_state.words:
    for i, word in enumerate(st.session_state.words):
        st.write(f"Annotate the word: {word}")
        st.session_state.labels[i] = st.text_input(
            f"Label for '{word}'", 
            value=st.session_state.labels[i], 
            key=f"label_{i}"
        )
        st.session_state.colors[i] = st.color_picker(
            f"Color for '{word}'", 
            value=st.session_state.colors[i], 
            key=f"color_{i}"
        )

    # Generate button to process the annotations
    if st.button("Generate Annotated Text", type="primary"):
        annotated_elements = []
        for i, word in enumerate(st.session_state.words):
            if st.session_state.labels[i] and st.session_state.colors[i] != "#FFFFFF":
                annotated_elements.append((word, st.session_state.labels[i], st.session_state.colors[i]))
            else:
                annotated_elements.append(word)
            annotated_elements.append(" ")  # Add space between words

        # Remove the last extra space added
        if annotated_elements and annotated_elements[-1] == " ":
            annotated_elements.pop()

        # Display the annotated text using the `annotated_text` function
        st.subheader("Annotated Text:")
        annotated_text(*annotated_elements)

        # Print the code for the annotated text
        st.subheader("Generated Code:")
        code_str = 'annotated_text(\n'
        for elem in annotated_elements:
            if isinstance(elem, tuple):
                code_str += f'    ("{elem[0]}", "{elem[1]}", "{elem[2]}"),\n'
            else:
                code_str += f'    "{elem}",\n'
        code_str += ')'
        st.code(code_str, language='python')