import gradio as gr
import random
from huggingface_hub import InferenceClient
# SEMANTIC SEARCH STEP 1
from sentence_transformers import SentenceTransformer
import torch

# SEMANTIC SEARCH STEP 2 --> EDIT WITH YOUR OWN KNOWLEDGE BASE WHEN READY
with open("skin_cancer_harvard.txt", "r", encoding="utf-8") as file:
  # Read the entire contents of the file and store it in a variable
  skin_cancer_harvard_text = file.read()
print(skin_cancer_harvard_text)

# SEMANTIC SEARCH STEP 3
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 newline character (\n)
  chunks = cleaned_text.split("\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:
    stripped_chunk = chunk.strip()
    cleaned_chunks.append(stripped_chunk)
  print(cleaned_chunks)

  print(len(cleaned_chunks))

  # Return the cleaned_chunks
  return cleaned_chunks

# Call the preprocess_text function and store the result in a cleaned_chunks variable
cleaned_chunks = preprocess_text(skin_cancer_harvard_text) # Complete this line; edit with my knowledgebase when ready

# SEMANTIC SEARCH STEP 4
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 text_chunks list

  # Print the chunk embeddings
  print(chunk_embeddings)


  # Print the shape of chunk_embeddings
  print(chunk_embeddings.shape)


  # 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

# SEMANTIC SEARCH STEP 5
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; bring it to the length of 1
  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

  # Print the similarities
  print(similarities)

  # Find the indices of the 3 chunks with highest similarity scores
  top_indices = torch.topk(similarities, k=3).indices
  print(top_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:
    relevant_info = cleaned_chunks[i]
    top_chunks.append(relevant_info)

  # Return the list of most relevant chunks
  return top_chunks

# SEMANTIC SEARCH STEP 6

# Call the get_top_chunks function with the original query
top_results = get_top_chunks('Is water good?',chunk_embeddings, cleaned_chunks) # Complete this line

print(top_results)# Print the top results

#the og code from gen ai lesson
client = InferenceClient("microsoft/phi-4")
# name of llm chatbot accessed ^^ or can use ' microsoft/phi-4 that's connected to the microsoft phi gen model

def respond(message,history):
    
    info = get_top_chunks(message, chunk_embeddings, cleaned_chunks)
    messages = [{'role': 'system','content':f'You are a friendly chatbot using {info} to answer questions.'}]
    #use string interporlation with variable info

    if history:
        messages.extend(history)

    messages.append({'role': 'user','content': message})

    response = client.chat_completion(messages, max_tokens = 500, top_p=0.8)
    #max tokens is a parameter to determine how long the message should be

    return response['choices'][0]['message']['content'].strip()


chatbot = gr.ChatInterface(respond, type='messages')
#defining my chatbot so user can interact, see their conversation and send new messages

chatbot.launch()