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) 
  # 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) 


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()

# 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
  # 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#
url = "https://mentalhealthfirstaid.org/mental-health-resources/"
yt = "https://youtu.be/7CCTOvZH0KU?si=6G80QeaA4cKssFeX"


about_text = f"""
<h2>About this bot</h2>
<p>Welcome to Mind Matters, an online resource that reminds
<em>you that your mind matters</em></p>

<p>Disclaimer: Mind Matters should not be used as an
alternative to seeking professional help. I am simply a
support tool.</p>
<p>All credits to the owner of the videos. We do not own any of the resources provided. 

<p>Click <a href="{url}" target="_blank">Free Resources</a> to access Free Mental Health Resources.</p>
<p>Click <a href="{yt}" target="_blank">here</a> to learn more about mental health.</p>
<p>You've got this! :) .</p>
"""

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=1):
            gr.HTML(about_text)
        with gr.Column(scale=2):
            gr.ChatInterface(fn=respond, title = "Mind Matters", description = "Always here to help", editable = True)

#demo.launch(theme=gr.themes.Soft().set(body_background_fill = "#20235c", body_text_color = "#c7a50e", block_background_fill = "b9d3eb", block_border_color = "#FFD700", button_primary_text_color = "#dae9f7")) 
 
demo.launch(
   theme=gr.themes.Soft().set(
        body_background_fill="#7F9CBB",
        body_text_color="black",
        block_background_fill="#49698D",
        block_border_color="#33455B",
        button_primary_background_fill="#49698D",
        button_primary_text_color="white"
    )
)