import gradio as gr
from sentence_transformers import SentenceTransformer
from huggingface_hub import InferenceClient
import pandas as pd
import torch
import math
import httpcore
import pickle
import time
setattr(httpcore, 'SyncHTTPTransport', 'AsyncHTTPProxy')

client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
model = SentenceTransformer('intfloat/multilingual-e5-large-instruct')

def get_detailed_instruct(task_description: str, query: str) -> str:
        return f'Instruct: {task_description}\nQuery: {query}'

def respond(message,
    max_tokens = 2048,
    temperature = 0.7,
    top_p = 0.95,
    ):
    #system role
    messages = [{"role": "system", "content": "You are a sunni moslem bot that always give answer based on quran, hadith, and the companions of prophet Muhammad!"}]

    #make a moslem bot
    messages.append({"role": "user", "content": "I want you to answer strictly based on quran and hadith"})
    messages.append({"role": "assistant", "content": "I'd be happy to help! Please go ahead and provide the sentence you'd like me to analyze. Please specify whether you're referencing a particular verse or hadith (Prophetic tradition) from the Quran or Hadith, or if you're asking me to analyze a general statement."})

    #adding fatwa references
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    selected_references = torch.load('selected_references.sav', map_location=torch.device(device))
    encoded_questions = torch.load('encoded_questions.sav', map_location=torch.device(device))
    
    task = 'Given a web search query, retrieve relevant passages that answer the query'
    queries = [
        get_detailed_instruct(task, message)
    ]
    print("start\n")
    print(time.time())

    query_embeddings = model.encode(queries, convert_to_tensor=True, normalize_embeddings=True)
    scores = (query_embeddings @ encoded_questions.T) * 100
    selected_references['similarity'] = scores.tolist()[0]
    sorted_references = selected_references.sort_values(by='similarity', ascending=False)
    sorted_references = sorted_references.iloc[:1]
    sorted_references = sorted_references.sort_values(by='similarity', ascending=True)
    print(sorted_references.shape[0])
    print(sorted_references['similarity'].tolist())
    print("sorted references\n")
    print(time.time())

    from googletrans import Translator
    translator = Translator()
    
    for index, row in sorted_references.iterrows():
        if(type(row["user"]) is str and type(row['assistant']) is str):
            try:
                translator = Translator()
                print(index)
                print(f'{row["user"]}')
                translated = translator.translate(f'{row["user"]}', src='ar', dest='en')
                print(translated)
                user = translated.text
                print(user)
                assistant = translator.translate(row['assistant']).text
                messages.append({"role": "user", "content":user })
                messages.append({"role": "assistant", "content": assistant})
            except Exception as error:
                print("An error occurred:", error)
                print("adding fatwa references exception occurred")
    
    print("append references\n")
    print(time.time())
    
    #adding more references
    df = pd.read_csv("moslem-bot-reference.csv", sep='|')
    for index, row in df.iterrows():
        messages.append({"role": "user", "content": row['user']})
        messages.append({"role": "assistant", "content": row['assistant']})
    
    print("added more references\n")
    print(time.time())
    
    #history from chat session
    """
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})
    """

    #latest user question
    from googletrans import Translator
    translator = Translator()
    en_message = ""
    message_language = "en"
    print("===message===")
    print(message)
    print("============")
    try:
        translator = Translator()
        print(translator.detect(message))
        message_language = translator.detect(message).lang
        print(message_language)
        print(translator.translate(message))
        en_message = translator.translate(message).text
        messages.append({"role": "user", "content": en_message})
    except Exception as error:
        messages.append({"role": "user", "content": message})
        print("An error occurred:", error)
        print("en_message exception occurred")
    
    #print(messages)
    print("added last question\n")
    print(time.time())

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        try:
            token = message.choices[0].delta.content
            response += token
            translated_response = translator.translate(response, src='en', dest=message_language).text
            yield translated_response
        except Exception as error:
            print("An error occurred:", error)
            yield ""
    
demo = gr.Interface(
    fn=respond,
    additional_inputs=[
        gr.Slider(minimum=1, maximum=2048, value=2048, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ], 
    inputs="textbox", 
    outputs="textbox",  
    cache_examples="lazy",
    examples=[
                ["Why is men created?"],
                ["Please tell me about superstition!"],
                ["How moses defeat pharaoh?"],
            ],
    title="Moslem Bot")
    
if __name__ == "__main__":
    demo.launch()