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

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

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

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

    cleaned_chunks = []

    for chunk in chunks:
        cleaned_chunk = chunk.strip()

        if (cleaned_chunk != ""):
            cleaned_chunks.append(cleaned_chunk)

    return cleaned_chunks

def create_embeddings(text_chunks):
    chunk_embeddings = model.encode(text_chunks, convert_to_tensor=True)
    return chunk_embeddings

def get_top_chunks(query, chunk_embeddings, text_chunks):
    query_embedding = model.encode(query, convert_to_tensor=True)
    
    query_embedding_normalized = query_embedding / query_embedding.norm()

    chunk_embeddings_normalized = chunk_embeddings / chunk_embeddings.norm(dim=1, keepdim=True)

    similarities = torch.matmul(chunk_embeddings_normalized, query_embedding_normalized)
    
    top_indices = torch.topk(similarities, k=8).indices

    top_chunks = []

    for i in top_indices:
        top_chunks.append(text_chunks[i])

    return top_chunks

def respond(message, history):
    top_results = get_top_chunks(message, chunk_embeddings, cleaned_chunks)
    
    messages = [{"role": "system",
                "content": 
                 """You are a friendly cafe recomender who aims to recomend Seattle area cafes. 
                Specifically, you have a lot of knowledge about Seattle cafe pros and cons, whether
                 they're good for working/studying, whether they're good for what group sizes, and more.
                 You should focus on recommending cafes in Seattle, not other activities/restaurants. Bakeries
                 can also be recommended. For tourists, don't prioritize cafes that are good for working/studying since
                 tourists likely aren't in town to do those things. Clearly describe the pros and cons of each cafe suggested using bullet points
                 of short sentences in the following format:
                Cafe Name (bolded) (Address) : Short one-liner describing the cafe (line break)
                Pros: 
                    - bullet points below of short sentences
                Cons:
                    - bullet points below of short sentences
                Repeat the above format for each cafe in your response

                If asked to plan an itinerary, ensure that all locations suggested are in the area requested/are close to each other. Use the following
                format when a user requests an itinerary:
                Stop (Number): Cafe Name (bolded) (Address): Short one-liner describing the cafe (line break)
                Pros:
                    - bullet points below of short sentences
                Cons:
                    - bullet pionts below of short sentences
                Repeat the above format for each cafe in the itinerary"""},
               {"role": "system",
               "content": f"Use the following cafes in your response in the order given: {top_results} If any clearly match the question, reccomend it first"}]
    response = ""
    
    if history:
        messages.extend(history)
        
    messages.append({"role": "user", "content": message})

    for message in client.chat_completion(
        messages,
        max_tokens = 1500,
        stream = True,
        temperature = 0.8
    ):
        if message.choices:
            token = message.choices[0].delta.content
        
        if token is not None:
            response += token
                
        yield response

cleaned_chunks = preprocess_text(knowledge_text)
chunk_embeddings = create_embeddings(cleaned_chunks)

with gr.Blocks() as interface:
    gr.Image(
	    value = "crema-banner.png",
        buttons = [],
        show_label = False)
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("""
            ## About Crema
            Crema can help find you a personalized Seattle cafe reccomendation whether you're looking for somewhere to work,
            somewhere to catch up with friends, or anything in between! Simply type a question to the chatbot or choose one of the
            exmaple questions to get started ☕🌱
            """)

        with gr.Column(scale=2):
            gr.ChatInterface(respond, title = "Crema", examples = ["What are some good cafes for studying?", "What cafes do I need to go to as a tourist?"])

interface.launch(theme = gr.Theme.from_hub("kbray/NeoSand"), share = True)