import streamlit as st
from sentence_transformers import SentenceTransformer
from qdrant_client import models, QdrantClient
import pandas as pd
from datasets import load_dataset

#************************************************************* LOAD DATA
data = load_dataset("ManuelAlv/academic_conuseling")

# Main dataset
bert_dataset = data['dataset'].to_pandas()

# Dataset used to test the chatbot
test_dataset = data['test'].to_pandas()
test_dataset.columns = ["test_question", "original_question"]

#************************************************************* Create Functions
# function to add values
def add_value(collection, key, value, id):
  encoder = SentenceTransformer(collection)

  qdrant.upsert(
      collection_name = collection,
      wait=True,
      points = [
          models.PointStruct(
              id = id,
              vector = encoder.encode(key).tolist(),
              payload = {
                  'text': value,
                  'question': key
              }
          )
      ]
  )

# Function to search for a value
def search(collection, query):
  search = qdrant.search(
      collection_name = collection,
      query_vector = encoder.encode(query).tolist(),
      limit = 1
  )
  return search

#************************************************************* Create VD
# Create a local Vector Database
qdrant = QdrantClient(":memory:")

# Load the model
model = "all-MiniLM-L6-v2"
encoder = SentenceTransformer(model)

# Create a collection for the model with its embeddings
qdrant.recreate_collection(
    collection_name = model,
    vectors_config = models.VectorParams(
        size = encoder.get_sentence_embedding_dimension(),
        distance = models.Distance.COSINE
    )
)

# Add the data to model
for index, row in bert_dataset.iterrows():
  key = row['question']
  value = row['answer']
  id = index + 1

  add_value(model, key, value, id)

# ************************************************************* QUERY
# Enter a question
question = "I'm feeling sad and lonely"

result = search(model, question)
result = result[0].payload['text']


def get_response(input):
    result = search(model, input)
    result = result[0].payload['text']
    return result

st.set_page_config(page_title="RAG", page_icon="🧊", layout="wide")
st.title("UniSA Academic Support")

# with st.sidebar:
#     st.header("Settings")
#     st.text_input("Enter a website URL")

if 'conversation_ended' not in st.session_state:
    st.session_state['conversation_ended'] = False

if not st.session_state['conversation_ended']:
  with st.chat_message("AI"):
    st.write("Hi! I'm BrainHug AI, your supportive AI friend.")
    st.write("Feel free to chat with me at any time, just enter your question. If I can't answer, try rephrasing it again.")
    st.write("If you want to finish the conversation, just say BYE")

  user_q = st.chat_input("Start typing here")

  if user_q:
    if user_q.upper() == "BYE":
      st.session_state['conversation_ended'] = True
      with st.chat_message("AI"):
        st.write("Goodbye! Feel free to come back anytime.")
      st.stop()

    elif user_q is not None or user_q is "":
      response = get_response(user_q)
      with st.chat_message("Human"):
          st.write(user_q)
      with st.chat_message("AI"):
          st.write(response)
else:
    st.write("The conversation has ended. Refresh the page to start over.")