import os
import shutil
import gradio as gr
import pandas as pd
from pandasai import PandasAI
from pandasai.llm.huggingface import HuggingFaceLLM
from transformers import HfEngine
from gradio import Chatbot
from huggingface_hub import login
from gradio.data_classes import FileData

# Log in to Hugging Face
login(os.getenv("HUGGINGFACEHUB_API_TOKEN"))

# Initialize the LLM engine with an Arabic model
llm_engine = HfEngine("Rohan-Kurdekar/Arabic_Bert_Model")

# Initialize the agent with PandasAI
llm = HuggingFaceLLM(engine=llm_engine)
pandas_ai = PandasAI(llm)

# Define the base prompt
base_prompt = """أنت محلل بيانات خبير.
وفقًا للميزات التي لديك وهيكل البيانات الموضح أدناه، حدد أي ميزة يجب أن تكون الهدف.
ثم قم بإدراج 3 أسئلة مثيرة للاهتمام يمكن طرحها على هذه البيانات، على سبيل المثال حول الارتباطات المحددة مع متغير الهدف.
ثم أجب عن هذه الأسئلة واحدًا تلو الآخر، من خلال العثور على الأرقام ذات الصلة.
وفي الوقت نفسه، قم برسم بعض الرسوم البيانية باستخدام matplotlib/seaborn واحفظها في المجلد './figures/' (الموجود مسبقًا): تأكد من مسح كل رسم بياني باستخدام plt.clf() قبل القيام برسم بياني آخر.
في إجابتك النهائية: قم بتلخيص هذه الارتباطات والاتجاهات
بعد كل رقم اشتق رؤى من العالم الحقيقي، على سبيل المثال: "الارتباط بين is_december والشعور بالملل هو 1.3453، مما يشير إلى أن الناس يشعرون بالملل أكثر في الشتاء".
يجب أن تكون إجابتك النهائية سلسلة طويلة تحتوي على 3 أجزاء مفصلة ومرقمة على الأقل.
هيكل البيانات:
{structure_notes}
يتم تمرير ملف البيانات إليك كإطار بيانات pandas، يمكنك استخدامه مباشرة.
لا تحاول تحميل data_file، فهي بالفعل إطار بيانات محمل مسبقًا في مترجم Python الخاص بك!
"""

# Example notes for Arabic dataset
example_notes = """هذه البيانات تتعلق بنتائج الطلاب.
الهدف من هذه البيانات هو معرفة الدرجة النهائية للطلاب.
name: اسم الطالب
age: عمر الطالب
gender: جنس الطالب
score: درجة الطالب"""

def interact_with_agent(file_input, additional_notes):
    try:
        # Ensure the ./figures directory exists
        figures_dir = "./figures"
        if os.path.exists(figures_dir):
            shutil.rmtree(figures_dir)
        os.makedirs(figures_dir)

        # Read the CSV file with the appropriate encoding for Arabic text
        data_file = pd.read_csv(file_input.name, encoding='ISO-8859-1')

        # Use PandasAI to analyze the data
        prompt = base_prompt.format(structure_notes="")
        if additional_notes and len(additional_notes) > 0:
            prompt += "\nملاحظات إضافية عن البيانات:\n" + additional_notes

        response = pandas_ai.run(data_file, prompt=prompt)
        
        messages = [gr.ChatMessage(role="user", content=prompt)]
        messages.append(gr.ChatMessage(role="assistant", content=response))

        yield messages

    except Exception as e:
        print(f"Error: {e}")
        yield [gr.ChatMessage(role="assistant", content=f"Error: {e}")]

with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue=gr.themes.colors.yellow,
        secondary_hue=gr.themes.colors.blue,
    )
) as demo:
    gr.Markdown("""# محلل البيانات Llama-3.1 📊🤔
قم بإسقاط ملف `.csv` أدناه، وأضف ملاحظات لوصف هذه البيانات إذا لزم الأمر، وسيقوم **Llama-3.1-70B بتحليل محتوى الملف ورسم الأشكال من أجلك!**""")
    file_input = gr.File(label="ملفك للتحليل")
    text_input = gr.Textbox(
        label="ملاحظات إضافية لدعم التحليل"
    )
    submit = gr.Button("تشغيل التحليل!", variant="primary")
    chatbot = gr.Chatbot(
        label="وكيل محلل البيانات",
        type="messages",
        avatar_images=(
            None,
            "https://em-content.zobj.net/source/twitter/53/robot-face_1f916.png",
        ),
    )
    gr.Examples(
        examples=[["./example/a.csv", example_notes]],
        inputs=[file_input, text_input],
        cache_examples=False
    )

    submit.click(interact_with_agent, [file_input, text_input], [chatbot])

if __name__ == "__main__":
    demo.launch(share=True)