import datahorse
import streamlit as st
import pandas as pd
import matplotlib.pyplot as plt
import io
import sys

st.title('Text to Data Analysis')

# Initialize session state for conversation history
if 'conversation' not in st.session_state:
    st.session_state.conversation = []

uploaded_file = st.file_uploader('Upload .csv data: ')


query = st.text_input('Search about data: ')

def capture_output(func):
    # Capture stdout
    old_stdout = sys.stdout
    new_stdout = io.StringIO()
    sys.stdout = new_stdout
    
    # Capture matplotlib figures
    fig = plt.figure()
    
    result = func()
    
    # Reset stdout
    sys.stdout = old_stdout
    output = new_stdout.getvalue()
    
    # Check if a plot was created
    if plt.gcf().axes:
        return fig
    elif output:
        return output
    else:
        return result

if uploaded_file is not None:

    # Add some example queries
    st.sidebar.header("Example Queries")
    st.sidebar.write("1. Show me a summary of the data")
    st.sidebar.write("2. Create a bar chart of [column_name]")
    st.sidebar.write("3. What is the average of [column_name]?")
    st.sidebar.write("4. Show me the correlation between [column1] and [column2]")
    st.sidebar.write("5. List the unique values in [column_name]")
    st.sidebar.write("6. What is the maximum value in [column_name]?")
    st.sidebar.write("7. Create a line chart of [column_name] over time")
    st.sidebar.write("8. How many missing values are in [column_name]?")
    st.sidebar.write("9. Filter rows where [column_name] is greater than [value]")
    st.sidebar.write("10. Generate a pie chart for [column_name]")


    df = datahorse.read(uploaded_file)

    col1, col2 = st.columns(2)


    if col1.button('Search it'):
        # Append user query to conversation history
        user_message = f"""
        <div style='background-color: #f0f0f0; padding: 10px; border-radius: 10px; margin: 5px 0;'>
            <strong>Me:</strong> {query}
        </div>
        """
        st.session_state.conversation.append(user_message)
        
        # Get response from datahorse
        response = capture_output(lambda: df.chat(query))
        
        st.subheader("Response:")
        
        if isinstance(response, plt.Figure):
            st.pyplot(response)
        elif isinstance(response, str):
            st.text(response)
        elif isinstance(response, pd.DataFrame):
            st.dataframe(response)
        elif response is not None:
            st.write(response)
        else:
            st.write("No output was captured.")

        # Append response to conversation history
        response_message = f"""
        <div style='background-color: #4CAF50; color: white; padding: 10px; border-radius: 10px; margin: 5px 0;'>
            <strong>Datahorse:</strong> Response displayed above
        </div>
        """
        st.session_state.conversation.append(response_message)
    # Clear conversation history
    if col2.button('Clear Conversation'):
        st.session_state.conversation = []


# Display conversation history
for message in reversed(st.session_state.conversation):
    st.markdown(message, unsafe_allow_html=True)

# Display the dataframe
if uploaded_file is not None:
    st.write("Preview of the uploaded data:")
    st.dataframe(df.head())
    
    # Display column names
    st.write("Column names in the dataset:")
    st.write(", ".join(df.columns))