### importing required libraries
from pandasai.responses.streamlit_response import StreamlitResponse 
from langchain_groq import ChatGroq
from pandasai import SmartDataframe
import streamlit as st
import pandas as pd
import shutil
import os

### using Mixtral 8 x 7b model with Groq's LPU Inference Engine
llm = ChatGroq(model_name = "mixtral-8x7b-32768", temperature = 0.1)

### building the streamlit interface
st.title("Smart Analyzer: Natural Language, Powerful Results")
dataFile = st.file_uploader(label = "Upload a csv/xlsx file for analysis", type = ["csv", "xlsx"])

if dataFile != None:
    extension = dataFile.name.split(".")[-1]
    if extension.lower() == "csv":
        dataframe = pd.read_csv(dataFile)
    elif extension.lower() == "xlsx":
        dataframe = pd.read_excel(dataFile)
    else:   
        st.warning("Invalid File Chosen")

    st.write(dataframe.head(3))
    
    smartDataFrame = SmartDataframe(dataframe, config = {
        "llm": llm,
        "custom_whitelisted_dependencies": ["IPython", "os", "matplotlib", "seaborn", "pandas"],
        "response_parser": StreamlitResponse,
        "verbose": True
    }
    )

    prompt = st.text_area("Enter your query, question, or desired operation")

    if st.button("Generate Answer"):
        path = os.path.join("exports", "charts")
        if os.path.isdir(path):
            pass
        else:
            os.makedirs(path)
        chartsDirectory = path
        if os.path.isdir(chartsDirectory): 
            if len(os.listdir(chartsDirectory)) != 0:
                shutil.rmtree(chartsDirectory)
            else: pass
            os.rmdir(chartsDirectory)
        else: pass

        if prompt:
            with st.spinner("Analyzing data, generating a response..."):
                response = smartDataFrame.chat(prompt)
                if len(os.listdir(chartsDirectory)) != 0:
                    imgName = os.listdir(chartsDirectory)[0]
                    imgPath = os.path.join(chartsDirectory, imgName)
                    st.image(imgPath)
                    shutil.rmtree(chartsDirectory)
                else:    
                    st.write(response)
        else:
            st.warning("Please enter a prompt")