import gradio as gr
from huggingface_hub import InferenceClient
from sentence_transformers import SentenceTransformer
import torch

#create custom CSS
css = """
.border-box {
    border: 6px solid #390099 !important;
    padding: 10px !important;
    border-radius: 10px;
}
"""


# Open the water_cycle.txt file in read mode with UTF-8 encoding
with open("BookBuddy1.txt", "r", encoding="utf-8") as file:
  # Read the entire contents of the file and store it in a variable
  book_text_part1 = file.read()

with open("BookBuddy2.txt", "r", encoding="utf-8") as file:
  # Read the entire contents of the file and store it in a variable
  book_text_part2 = file.read()

book_text = book_text_part1 + "\n\n" + book_text_part2

print(book_text)


def preprocess_text(text):
  # Strip extra whitespace from the beginning and the end of the text
  cleaned_text = text.strip()

  # Split the cleaned_text by every period
  chunks = cleaned_text.split("\n\n")

  # Create an empty list to store cleaned chunks
  cleaned_chunks = []

  # Write your for-in loop below to clean each chunk and add it to the cleaned_chunks list
  for chunk in chunks:
    chunk.strip()

    if chunk != "":
      cleaned_chunks.append(chunk)
    
    cleaned_chunks = [str(c) for c in cleaned_chunks if c and len(str(c).strip()) > 0]

  # Return the cleaned_chunks
  return cleaned_chunks

model = SentenceTransformer('all-MiniLM-L6-v2')

def create_embeddings(text_chunks):
  # Convert each text chunk into a vector embedding and store as a tensor
  chunk_embeddings = model.encode(text_chunks, convert_to_tensor=True)

  # Return the chunk_embeddings
  return chunk_embeddings

def get_top_chunks(query, chunk_embeddings, text_chunks):
    # Convert the query text into a vector embedding
    query_embedding = model.encode(query, convert_to_tensor=True)
    
    # Normalize the query embedding to unit length for accurate similarity comparison
    query_embedding_normalized = query_embedding / query_embedding.norm()

    print(chunk_embeddings)
    # Normalize all chunk embeddings to unit length for consistent comparison
    chunk_embeddings_normalized = chunk_embeddings / chunk_embeddings.norm(dim=1, keepdim=True)
    
    # Calculate cosine similarity between query and all chunks using matrix multiplication
    similarities = torch.matmul(chunk_embeddings_normalized, query_embedding_normalized)
    
    # Find the indices of the 3 chunks with highest similarity scores
    top_indices = torch.topk(similarities, k=3).indices
    
    # Create an empty list to store the most relevant chunks
    top_chunks = []
    
    # Loop through the top indices and retrieve the corresponding text chunks
    for i in top_indices:
        top_chunks.append(text_chunks[i])
        
    # Return the list of most relevant chunks
    return top_chunks


cleaned_chunks = preprocess_text(book_text)
chunk_embeddings = create_embeddings(cleaned_chunks)



client = InferenceClient("meta-llama/Llama-3.1-8B-Instruct")

def respond(message, history):
    
    messages = [{"role": "system", "content": "You are a friendly chatbot who provides book recommendations including but not limited to age and interests, but you must ask for these features or they must be volunteered to you in the message. You also provide reasons for your book recommendation. The language you use in your responses should be accessible to the average 13-18 year old. Your name is BookBuddy."}]

    context = get_top_chunks(message, chunk_embeddings, cleaned_chunks)

    prompt = f"""Use the following information to answer:
    Context: {context}
    Question: {message}"""
    
    if history:
        messages.extend(history)
        
    messages.append({"role": "user", "content": prompt})

    response = ""

    for chunk in client.chat_completion(
        model="meta-llama/Llama-3.1-8B-Instruct",
        temperature=0.25,
        messages=messages,
        max_tokens=5000,
        stream=True
    ):
        if not chunk.choices:
            continue

        token = chunk.choices[0].delta.content

        if token:
            response += token
            yield response

about_text = "Hello! My name is BookBuddy and I'm here to help you find new and exciting books based on your preferences. Let's get reading 😁📖"

with gr.Blocks(theme=gr.Theme.from_hub("hmb/windows95"), css=css) as demo:
    with gr.Column():
        with gr.Row(scale=1, elem_classes="border-box"):
            gr.Markdown(about_text)

        with gr.Row(scale=2):
            gr.ChatInterface(respond)

demo.launch()