import gradio as gr
from huggingface_hub import InferenceClient
from sentence_transformers import SentenceTransformer, util
import torch
import pandas as pd

def load_bible():
    # Replace: 1=Genesis, 2=Exodus, ... 66=Revelation
    books = ['Genesis', 'Exodus', 'Leviticus', 'Numbers', 'Deuteronomy',
           'Joshua', 'Judges', 'Ruth', '1 Samuel', '2 Samuel', '1 Kings',
           '2 Kings', '1 Chronicles', '2 Chronicles', 'Ezra', 'Nehemiah',
           'Esther', 'Job', 'Psalms', 'Proverbs', 'Ecclesiastes',
           'Song of Solomon', 'Isaiah', 'Jeremiah', 'Lamentations', 'Ezekiel',
           'Daniel', 'Hosea', 'Joel', 'Amos', 'Obadiah', 'Jonah', 'Micah',
           'Nahum', 'Habakkuk', 'Zephaniah', 'Haggai', 'Zechariah', 'Malachi',
           'Matthew', 'Mark', 'Luke', 'John', 'Acts', 'Romans',
           '1 Corinthians', '2 Corinthians', 'Galatians', 'Ephesians',
           'Philippians', 'Colossians', '1 Thessalonians', '2 Thessalonians',
           '1 Timothy', '2 Timothy', 'Titus', 'Philemon', 'Hebrews', 'James',
           '1 Peter', '2 Peter', '1 John', '2 John', '3 John', 'Jude',
           'Revelation']
    lookup = {}
    for i, j in enumerate(books):
        lookup[i+1]=j
    lookup

    esv = pd.read_csv('t_esv.csv').\
        rename(columns = {"b":"book"}).drop(["id"], axis=1)
    esv.book = esv.book.replace(lookup)
    return esv

def load_embeddings(fn):
    with open(fn, "rb") as f:
        embeddings = torch.load(f)
    return embeddings

def search(searchText, k = 5, show_html=False, return_str=False):
    emb = model.encode(searchText, convert_to_tensor=True)
    cos_scores = util.cos_sim(emb, embeddings)[0]
    top_results = torch.topk(cos_scores, k = k)

    x = bible.iloc[top_results.indices.cpu().detach().numpy(), :]

    if show_html:
        for index, row in x.iterrows():
            display(HTML(f'<p style="font-size: 22px;"><strong>{row.book} {row.c}:{row.v}</strong>  <em>{row.t}</em></p><br>'))

    if return_str:
        s = []
        for index, row in x.iterrows():
            tmp = f"({row.book} {row.c}:{row.v}) {row.t}"
            s.append(tmp)

        return "\n".join(s)

    else:
        return x


def create_prompt(q):
    ctx = search(q, k=10, return_str = True)
    s = f"""Context:
    {ctx}
    Question:
    {q}
    """
    
    return s

# Constants
model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
embeddings = load_embeddings("esv_embeddings.pt")
bible = load_bible()

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""

client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
# client = InferenceClient("microsoft/Phi-3-mini-4k-instruct")

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    messages = [{"role": "system", "content": system_message}]
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    message = create_prompt(message)
    messages.append({"role": "user", "content": message})

    response = ""
    
    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content

        response += token
        yield response

"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
demo = gr.ChatInterface(
    fn=respond,
    chatbot = gr.Chatbot(value = [[None, "What can I help you with today, beloved?"]]),
    additional_inputs=[
        gr.Textbox(value="You are a Christian Pastor. Provide spiritual wisdom, based on biblical truth found in the context. Quote from the context when appropriate.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, 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)",
        ),
    ],
)


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