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

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

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

print(orton_knowledge_text)

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)

    print(cleaned_chunks)

    print(len(cleaned_chunks)) # including prints so I can check as I go

    return cleaned_chunks

cleaned_chunks = preprocess_text(orton_knowledge_text)

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

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

    print(chunk_embeddings)

    print(chunk_embeddings.shape)

    return chunk_embeddings

chunk_embeddings = create_embeddings(cleaned_chunks)

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)

    print(similarities)

    top_indices = torch.topk(similarities, k=3).indices

    print(top_indices)

    top_chunks = []

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

    return(top_chunks)




def respond(message, history):
    top_chunks = get_top_chunks(message, chunk_embeddings, cleaned_chunks)

    context = "\n\n".join(top_chunks)
    messages = [{"role": "system", 
    "content": f"You are a phonics instruction expert."
               f"You will ask the user their child's age and grade." 
               f"You will ask questions to find out their current proficieny."
               f"You will respond with one website reccomendation and one skill practice."
               f"Use the following knowledge to help answer:\n\n{context}"
    }]
    
    if history:
        messages.extend(history)
        
    messages.append({"role": "user", "content": message})
    
    response = client.chat_completion(
        messages,
        max_tokens= 900, 
        temperature = .2, 
        frequency_penalty = 1, 
        stream = True
    )
    response_text = ""
    for message in response:
        if not message.choices:
            continue
        token = message.choices[0].delta.content
        if token is None:
            continue
        response_text += token
        yield response_text

    
chatbot = gr.ChatInterface(respond, title = "At Home Phonics Support", description ="Tell me your child's age and grade and I will recommend at home supports.")
   
chatbot.launch(share=True, debug=True)