import os
import torch
import gradio as gr
import pandas as pd
from sentence_transformers import SentenceTransformer, util
from transformers import GenerationConfig, AutoModelForCausalLM, AutoTokenizer

os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"


class PersianRAG:
    def __init__(self, knowledge, embedding_model='LABSE', llm_model="MehdiHosseiniMoghadam/AVA-Mistral-7B-V2",
                 device='cuda', retrieved_docs=3):
        self.device = device
        self.retrieved_docs = retrieved_docs
        self.answer_df = (knowledge['Answer'])

        self.embedder = SentenceTransformer(embedding_model)
        self.question_embeddings = self.embedder.encode((knowledge['Question']), show_progress_bar=True,
                                                        convert_to_tensor=True)

        self.model = AutoModelForCausalLM.from_pretrained(llm_model, torch_dtype=torch.float16, device_map="auto")
        self.tokenizer = AutoTokenizer.from_pretrained(llm_model)
        self.generation_config = GenerationConfig(
            do_sample=True,
            top_k=1,
            temperature=0.99,
            max_new_tokens=900,
            pad_token_id=self.tokenizer.eos_token_id
        )

    def rag(self, query):
        ans = {}
        question_embedding = self.embedder.encode(query, convert_to_tensor=True)
        hits = util.semantic_search(question_embedding, self.question_embeddings)
        hits = hits[0]

        for hit in hits[0:self.retrieved_docs]:
            ans[hit['corpus_id']] = self.answer_df[hit['corpus_id']]
        ans = pd.DataFrame(list(ans.items()), columns=['id', 'res'])

        prompt = f'''
        با توجه به شرایط زیر به این سوال پاسخ دهید:
        {query},

        متن نوشته:
        {ans['res'][0]} - {ans['res'][1]} - {ans['res'][2]}
        '''

        prompt = f"### Human:{prompt}\n### Assistant:"

        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
        outputs = self.model.generate(**inputs, generation_config=self.generation_config)

        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)


# Function to load CSV and initialize PersianRAG
def init_rag(knowledge_file, embedding_model, llm_model, device, retrieved_docs):
    knowledge = pd.read_csv(knowledge_file)
    rag_system = PersianRAG(knowledge, embedding_model=embedding_model, llm_model=llm_model, device=device,
                            retrieved_docs=retrieved_docs)
    return rag_system


# Function to handle querying
def query_rag(rag_system, query):
    return rag_system.rag(query)


# Gradio interface to upload CSV and configure RAG system
def rag_interface(knowledge_file, query, embedding_model, llm_model, device, retrieved_docs):
    rag_system = init_rag(knowledge_file, embedding_model, llm_model, device, retrieved_docs)
    return query_rag(rag_system, query)


# Create Gradio interface
interface = gr.Interface(
    fn=rag_interface,
    inputs=[
        gr.File(label="Upload Knowledge Base CSV"),
        gr.Textbox(label="Enter your query"),
        gr.Dropdown(choices=["LABSE", "paraphrase-multilingual-mpnet-base-v2"], value="LABSE", label="Embedding Model"),
        gr.Textbox(value="MehdiHosseiniMoghadam/AVA-Mistral-7B-V2", label="LLM Model Name"),
        gr.Dropdown(choices=["cuda", "cpu"], value="cuda", label="Device"),
        gr.Slider(minimum=1, maximum=5, step=1, value=3, label="Number of Retrieved Documents")],
    outputs="text",
    title="Persian RAG System",
    description="Upload a CSV file as the knowledge base, ask a question, and get an answer.")

# Launch the Gradio interface
interface.launch(debug=True)