import streamlit as st 
import pandas as pd
from pandasai import SmartDataframe
from pandasai.llm.openai import OpenAI

openai_api_key = st.secrets["openai_api_key"]

# create an LLM by instantiating OpenAI object, and passing API token
llm = OpenAI(api_token = openai_api_key)

st.title("Data analysis with PandasAI")

uploaded_file = st.file_uploader("Upload a CSV file for analysis", type=['csv'])

if uploaded_file is not None:
    df = pd.read_csv(uploaded_file)
    st.write(df.head(3))

    # create PandasAI object, passing the LLM   
    sdf = SmartDataframe(df, config={"llm": llm})

    prompt = st.text_area("Enter your question:")

    # Generate the answer
    if st.button("Find the answer"):
        if prompt:
            # call pandas_ai.run(), passing dataframe and prompt
            with st.spinner("Generating response..."):
                
                st.write("The answer is: ")
                st.write(sdf.chat(prompt))

                st.write("Here is the code to get the answer: ")
                code_string = sdf.last_code_generated

                # Split the string into separate lines
                code_lines = code_string.strip().split('\n')

                # Print the lines
                for line in code_lines:
                    st.write(line)

        else:
            st.warning("Please enter a question.")