import gradio as gr

from huggingface_hub import InferenceClient
import os
client = InferenceClient(model="Qwen/Qwen2.5-7B-Instruct", token=os.environ.get("HF"))
from sentence_transformers import SentenceTransformer
import torch

with open("knowledge.txt", "r", encoding="utf-8") as file:
    knowledge_text = file.read()

def preprocess_text(text):
    cleaned_text = text.strip()
    chunks = cleaned_text.split("\n")
    cleaned_chunks = []
    for chunk in chunks:
        stripped_chunk = chunk.strip()
        if len(stripped_chunk) > 0:
            cleaned_chunks.append(stripped_chunk)
    return cleaned_chunks
cleaned_chunks = preprocess_text(knowledge_text)

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) # Replace ... with the cleaned_chunks list
  # Return the chunk_embeddings
  return chunk_embeddings

# Call the create_embeddings function and store the result in a new chunk_embeddings variable
chunk_embeddings = create_embeddings(cleaned_chunks) #complete this line


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) # Complete this line

  # Normalize the query embedding to unit length for accurate similarity comparison
  query_embedding_normalized = query_embedding / query_embedding.norm()

  # 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) # Complete this line

  # 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
  # This is only one way scholars may write this, but there are other ways!
  for i in top_indices:
    chunk = text_chunks[i]
    top_chunks.append(chunk)


  # Return the list of most relevant chunks
  return top_chunks



def respond(message, history):
    messages = [{"role": "system",
                 "content":"You are an emotional support chatbot. You would not take about anything else other than mental health and helping the users. You need to make sure the user is comfortable."
                }] 
    if history:
        messages.extend(history)
    messages.append({"role":"user",
                 "content":message
                })
    response = " "
    for msg in client.chat_completion(messages, max_tokens = 1000, temperature = 1, top_p = 0.5, stream = True):
        token = msg.choices[0].delta.content
        response += token 
        yield response 


#EMMA'S PRACTICE EDITS#
about_text = """
## About this bot  
    Welcome to Mind Matters, an online resource that reminds *your that your mind matters*
    Disclaimer; Mind Matters should not be used as an alternative to seeking professional help. I am simply a support tool.
"""

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown(about_text)

        with gr.Column(scale=2):
            gr.ChatInterface(fn=respond, title = "Mind Matters", description = "Always here to help", editable = True)

demo.launch()


chatbot = gr.ChatInterface(fn=respond, title = "Mind Matters", description = "Always here to help", editable = True)
chatbot.launch()