import os
import json
import re
from sentence_transformers import SentenceTransformer, CrossEncoder
import hnswlib
import numpy as np
from typing import Iterator

import gradio as gr
from gradio_client import Client
import pandas as pd
import torch

from transformers import AutoTokenizer
# from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer


MAX_MAX_NEW_TOKENS = 250
DEFAULT_MAX_NEW_TOKENS = 250
MAX_INPUT_TOKEN_LENGTH = 4000
EMBED_DIM = 1024
K = 2
EF = 100
TEXT_FILE = 'data.txt'
SEARCH_INDEX = "search_index.bin"
EMBEDDINGS_FILE = "embeddings.npy"
DOCUMENT_DATASET = "chunked_data.parquet"
COSINE_THRESHOLD = 0.7


torch_device = "cuda" if torch.cuda.is_available() else "cpu"
print("Running on device:", torch_device)
print("CPU threads:", torch.get_num_threads())

biencoder = SentenceTransformer("intfloat/e5-large-v2", device="cpu")
cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-12-v2", max_length=512, device="cpu")
model_name_or_path = "TheBloke/TinyLlama-1.1B-1T-OpenOrca-AWQ"

# Load model
# model = AutoAWQForCausalLM.from_quantized(model_name_or_path, fuse_layers=True,
#                                           trust_remote_code=False, safetensors=True)
# tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=False)
tokenizer = AutoTokenizer.from_pretrained("mosaicml/mpt-30b-chat", trust_remote_code=False)
chat_client = Client("https://mosaicml-mpt-30b-chat.hf.space/", serialize = False)
chat_bot = [["", None]]

def read_text_from_file(file_path):
    with open(file_path, "r", encoding="utf-8") as text_file:
        text = text_file.read()
    texts = text.split("&&")
    return [t.strip() for t in texts]

def create_qa_prompt(query, relevant_chunks):
    stuffed_context = " ".join(relevant_chunks)
    return f"""\
Use the following pieces of context given in to answer the question at the end. \
If you don't know the answer, just say that you don't know, don't try to make up an answer. \
Keep the answer short and succinct.
        
Context: {stuffed_context}
Question: {query}
Helpful Answer: \
"""


def create_condense_question_prompt(question, chat_history):
    return f"""\
Given the following conversation and a follow up question, \
rephrase the follow up question to be a standalone question in its original language. \
Output the json object with single field `question` and value being the rephrased standalone question. 
Only output json object and nothing else.
Chat History:
{chat_history}
Follow Up Input: {question}
"""


def get_prompt(message: str, chat_history: list[tuple[str, str]], system_prompt: str) -> str:
    texts = [f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n"]
    # The first user input is _not_ stripped
    do_strip = False
    for user_input, response in chat_history:
        user_input = user_input.strip() if do_strip else user_input
        do_strip = True
        texts.append(f"{user_input} [/INST] {response.strip()} </s><s>[INST] ")
    message = message.strip() if do_strip else message
    texts.append(f"{message} [/INST]")
    return "".join(texts)


def get_input_token_length(message: str, chat_history: list[tuple[str, str]], system_prompt: str) -> int:
    prompt = get_prompt(message, chat_history, system_prompt)
    input_ids = tokenizer([prompt], return_tensors="np", add_special_tokens=False)["input_ids"]
    return input_ids.shape[-1]

def prompt_builder(prompt, system_message="You are a helpful chatbot which gives correct and truthful answers"):
  return f'''<|im_start|>system
{system_message}<|im_end|>
<|im_start|>user
{prompt}<|im_end|>
<|im_start|>assistant


  '''
def get_completion(
    prompt,
    system_prompt=None,
    # model=model,
    max_new_tokens=250,
    temperature=0.2,
    top_p=0.95,
    top_k=50,
    stream=False,
    debug=False,
):
    global chat_bot
    if temperature < 1e-2:
        temperature = 1e-2
    answer=chat_client.predict(
                prompt, # str  in 'Type an input and press Enter' Textbox component
                chat_bot,
                fn_index=1
    )
    chat_bot = answer[1]
    yield answer[1][0][1]
    # prompt = prompt_builder(prompt)
    # tokens = tokenizer(
    #     prompt,
    #     return_tensors='pt'
    # ).input_ids.cuda()

    # # Generate output
    # for i in range(max_new_tokens):
    #   generation_output = model.generate(
    #       tokens,
    #       do_sample=True,
    #       temperature=temperature,
    #       top_p=top_p,
    #       top_k=top_k,
    #       max_new_tokens=1
    #   )
    #   tokens = generation_output
    #   yield tokenizer.decode(generation_output[0][-1])

# load the index for the data
def load_hnsw_index(index_file):
    # Load the HNSW index from the specified file
    index = hnswlib.Index(space="ip", dim=EMBED_DIM)
    index.load_index(index_file)
    return index


# create the index for the data from numpy embeddings
# avoid the arch mismatches when creating search index
def create_hnsw_index(embeddings_file, M=16, efC=100):
    embeddings = np.load(embeddings_file)
    # Create the HNSW index
    num_dim = embeddings.shape[1]
    ids = np.arange(embeddings.shape[0])
    index = hnswlib.Index(space="ip", dim=num_dim)
    index.init_index(max_elements=embeddings.shape[0], ef_construction=efC, M=M)
    index.add_items(embeddings, ids)
    return index


def create_query_embedding(query):
    # Encode the query to get its embedding
    embedding = biencoder.encode([query], normalize_embeddings=True)[0]
    return embedding


def find_nearest_neighbors(query_embedding):
    search_index.set_ef(EF)
    # Find the k-nearest neighbors for the query embedding
    labels, distances = search_index.knn_query(query_embedding, k=K)
    labels = [label for label, distance in zip(labels[0], distances[0]) if (1 - distance) >= COSINE_THRESHOLD]
    relevant_chunks = data_df.iloc[labels]["chunk_content"].tolist()
    return relevant_chunks


def rerank_chunks_with_cross_encoder(query, chunks):
    # Create a list of tuples, each containing a query-chunk pair
    pairs = [(query, chunk) for chunk in chunks]

    # Get scores for each query-chunk pair using the cross encoder
    scores = cross_encoder.predict(pairs)

    # Sort the chunks based on their scores in descending order
    sorted_chunks = [chunk for _, chunk in sorted(zip(scores, chunks), reverse=True)]

    return sorted_chunks


def generate_condensed_query(query, history):
    chat_history = ""
    for turn in history:
        chat_history += f"Human: {turn[0]}\n"
        chat_history += f"Assistant: {turn[1]}\n"

    condense_question_prompt = create_condense_question_prompt(query, chat_history)
    # condensed_question = json.loads(get_completion(condense_question_prompt, max_new_tokens=64, temperature=0))
    condensed_question = "".join([token for token in get_completion(condense_question_prompt, max_new_tokens=64, temperature=0)])
    return condensed_question


DEFAULT_SYSTEM_PROMPT = """\
You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\
"""
# MAX_MAX_NEW_TOKENS = 2048
# DEFAULT_MAX_NEW_TOKENS = 1024
# MAX_INPUT_TOKEN_LENGTH = 4000

DESCRIPTION = """
# AVA Southampton Chatbot 🤗
"""

LICENSE = """
<p/>
---

"""

if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>Running on CPU 🥶.</p>"


def clear_and_save_textbox(message: str) -> tuple[str, str]:
    return "", message


def display_input(message: str, history: list[tuple[str, str]]) -> list[tuple[str, str]]:
    history.append((message, ""))
    return history


def delete_prev_fn(history: list[tuple[str, str]]) -> tuple[list[tuple[str, str]], str]:
    try:
        message, _ = history.pop()
    except IndexError:
        message = ""
    return history, message or ""


def wrap_html_code(text):
    pattern = r"<.*?>"
    matches = re.findall(pattern, text)
    if len(matches) > 0:
        return f"```{text}```"
    else:
        return text


def generate(
    message: str,
    history_with_input: list[tuple[str, str]],
    system_prompt: str,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    top_k: int,
) -> Iterator[list[tuple[str, str]]]:
    if max_new_tokens > MAX_MAX_NEW_TOKENS:
        raise ValueError
    history = history_with_input[:-1]
    if len(history) > 0:
        condensed_query = generate_condensed_query(message, history)
        print(f"{condensed_query=}")
    else:
        condensed_query = message
    query_embedding = create_query_embedding(condensed_query)
    relevant_chunks = find_nearest_neighbors(query_embedding)
    reranked_relevant_chunks = rerank_chunks_with_cross_encoder(condensed_query, relevant_chunks)
    print(reranked_relevant_chunks)
    qa_prompt = create_qa_prompt(condensed_query, reranked_relevant_chunks)
    print(f"{qa_prompt=}")
    generator = get_completion(
        qa_prompt,
        system_prompt=system_prompt,
        #stream=True,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
    )

    output = ""
    for idx, response in enumerate(generator):
        # token = response["choices"][0]["delta"].get("content", "") or ""
        token = response
        output += token
        if idx == 0:
            history.append((message, output))
        else:
            history[-1] = (message, output)

        history = [
            (wrap_html_code(history[i][0]), wrap_html_code(history[i][1]))
            for i in range(0, len(history))
        ]
        yield history

    return history


def process_example(message: str) -> tuple[str, list[tuple[str, str]]]:
    generator = generate(message, [], DEFAULT_SYSTEM_PROMPT, 1024, 0.2, 0.95, 50)
    for x in generator:
        pass
    return "", x


def check_input_token_length(message: str, chat_history: list[tuple[str, str]], system_prompt: str) -> None:
    input_token_length = get_input_token_length(message, chat_history, system_prompt)
    if input_token_length > MAX_INPUT_TOKEN_LENGTH:
        raise gr.Error(
            f"The accumulated input is too long ({input_token_length} > {MAX_INPUT_TOKEN_LENGTH}). Clear your chat history and try again."
        )

if not os.path.exists(TEXT_FILE):
    os.system(f"wget -O {TEXT_FILE} https://huggingface.co/spaces/Slycat/Southampton-Similarity/resolve/main/Southampton.txt")

if not os.path.exists(EMBEDDINGS_FILE):
    texts = read_text_from_file(TEXT_FILE)
    embeddings = biencoder.encode(texts, normalize_embeddings=True)
    np.save(EMBEDDINGS_FILE,embeddings)

if not os.path.exists(DOCUMENT_DATASET):
    texts = read_text_from_file(TEXT_FILE)
    df = pd.DataFrame(texts, columns = ["chunk_content"])
    df.to_parquet(DOCUMENT_DATASET,index=False)

search_index = create_hnsw_index(EMBEDDINGS_FILE)  # load_hnsw_index(SEARCH_INDEX)
data_df = pd.read_parquet(DOCUMENT_DATASET).reset_index()
with gr.Blocks(css="style.css") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Group():
        chatbot = gr.Chatbot(label="Chatbot")
        with gr.Row():
            textbox = gr.Textbox(
                container=False,
                show_label=False,
                placeholder="Type a message...",
                scale=10,
            )
            submit_button = gr.Button("Submit", variant="primary", scale=1, min_width=0)
    with gr.Row():
        retry_button = gr.Button("🔄  Retry", variant="secondary")
        undo_button = gr.Button("↩️ Undo", variant="secondary")
        clear_button = gr.Button("🗑️  Clear", variant="secondary")

    saved_input = gr.State()

    with gr.Accordion(label="Advanced options", open=False):
        system_prompt = gr.Textbox(label="System prompt", value=DEFAULT_SYSTEM_PROMPT, lines=6)
        max_new_tokens = gr.Slider(
            label="Max new tokens",
            minimum=1,
            maximum=MAX_MAX_NEW_TOKENS,
            step=1,
            value=DEFAULT_MAX_NEW_TOKENS,
        )
        temperature = gr.Slider(
            label="Temperature",
            minimum=0.1,
            maximum=4.0,
            step=0.1,
            value=0.2,
        )
        top_p = gr.Slider(
            label="Top-p (nucleus sampling)",
            minimum=0.05,
            maximum=1.0,
            step=0.05,
            value=0.95,
        )
        top_k = gr.Slider(
            label="Top-k",
            minimum=1,
            maximum=1000,
            step=1,
            value=50,
        )

    gr.Examples(
        examples=[
            "What is University of Southampton?",
            "Is University of Southampton Good?",
            "What is sports facility at southampton university?",
            "How big is the Southampton campus?",
            "What are the rankings of southampton university?",
            "What research facilities does the Southampton university offer?"
        ],
        inputs=textbox,
        outputs=[textbox, chatbot],
        # fn=process_example,
        cache_examples=False,
    )

    gr.Markdown(LICENSE)

    textbox.submit(
        fn=clear_and_save_textbox,
        inputs=textbox,
        outputs=[textbox, saved_input],
        api_name=False,
        queue=False,
    ).then(fn=display_input, inputs=[saved_input, chatbot], outputs=chatbot, api_name=False, queue=False,).then(
        fn=check_input_token_length,
        inputs=[saved_input, chatbot, system_prompt],
        api_name=False,
        queue=False,
    ).success(
        fn=generate,
        inputs=[
            saved_input,
            chatbot,
            system_prompt,
            max_new_tokens,
            temperature,
            top_p,
            top_k,
        ],
        outputs=chatbot,
        api_name=False,
    )

    button_event_preprocess = (
        submit_button.click(
            fn=clear_and_save_textbox,
            inputs=textbox,
            outputs=[textbox, saved_input],
            api_name=False,
            queue=False,
        )
        .then(
            fn=display_input,
            inputs=[saved_input, chatbot],
            outputs=chatbot,
            api_name=False,
            queue=False,
        )
        .then(
            fn=check_input_token_length,
            inputs=[saved_input, chatbot, system_prompt],
            api_name=False,
            queue=False,
        )
        .success(
            fn=generate,
            inputs=[
                saved_input,
                chatbot,
                system_prompt,
                max_new_tokens,
                temperature,
                top_p,
                top_k,
            ],
            outputs=chatbot,
            api_name=False,
        )
    )

    retry_button.click(
        fn=delete_prev_fn,
        inputs=chatbot,
        outputs=[chatbot, saved_input],
        api_name=False,
        queue=False,
    ).then(fn=display_input, inputs=[saved_input, chatbot], outputs=chatbot, api_name=False, queue=False,).then(
        fn=generate,
        inputs=[
            saved_input,
            chatbot,
            system_prompt,
            max_new_tokens,
            temperature,
            top_p,
            top_k,
        ],
        outputs=chatbot,
        api_name=False,
    )

    undo_button.click(
        fn=delete_prev_fn,
        inputs=chatbot,
        outputs=[chatbot, saved_input],
        api_name=False,
        queue=False,
    ).then(
        fn=lambda x: x,
        inputs=[saved_input],
        outputs=textbox,
        api_name=False,
        queue=False,
    )

    clear_button.click(
        fn=lambda: ([], ""),
        outputs=[chatbot, saved_input],
        queue=False,
        api_name=False,
    )

demo.queue(max_size=20).launch(debug=True)