# +
import streamlit as st
from pandasai import SmartDataframe # requires 1.x numpy
from pandasai.llm import OpenAI
import matplotlib.pyplot as plt

import os

# +
import pandas as pd 
import seaborn as sns
columns = ['year', 'month', 'decimal_date', 'average', 'smooth', 'std_days', 'uncertainty_monthly_mean', 'empty']
df = pd.read_csv("https://gml.noaa.gov/webdata/ccgg/trends/co2/co2_mm_mlo.txt", sep="\s+", comment="#", names= columns)
df.to_csv("data.csv") 

st.session_state.df = df
st.session_state.prompt_history = []


# -

st.set_page_config(
    page_title="Chat tool",
    page_icon="🦜",
)
st.title("CO2 Explorer")

# +
# st.dataframe(df)
# -

llm = OpenAI(api_token=st.secrets["OPENAI_API_KEY"])
pandas_ai = SmartDataframe("data.csv", 
              config={
                      "llm": llm, 
                      "save_charts": False,
                      "verbose": False
                      }
            )


# +
chatbox = st.container()
with chatbox:  
    if question := st.chat_input("Plot the average CO2 concentration over time (decimal_date) using seaborn", key="chain"):
        st.chat_message("user").write(question)
        with st.chat_message("assistant"):
            x = pandas_ai.chat(question)
            if os.path.isfile(x):
                    im = plt.imread(x)
                    st.image(im)
                    os.remove(x)

            if x is not None:
                st.write(x)
            st.session_state.prompt_history.append(question)

    if st.session_state.df is not None:
        st.subheader("Current dataframe:")
        st.write(st.session_state.df)

    st.subheader("Prompt history:")
    st.write(st.session_state.prompt_history)


if st.button("Clear"):
    st.session_state.prompt_history = []
    st.session_state.df = None