# import required libraries
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.llms import HuggingFaceHub
from langchain_community.vectorstores import Chroma
from sentence_transformers import SentenceTransformer
from langchain_core.prompts import ChatPromptTemplate
from langchain import PromptTemplate

import streamlit as st
import sys,yaml
import chromadb
import Utilities as ut

hf_token=""
chromadbpath=""
chromadbcollname=""
embedding_model_id=""
llm_repo_id=""
#embeddings=None
#chroma_client=None


def filterdistance(distcoll):
    myemptydict={}
    if len(distcoll) < 0:myemptydict
    for distances in distcoll['distances']:
        for distance in distances:
            if distance<50: return distcoll
            else: return myemptydict
           
def get_collections(query):
    #myemptydict={}
    result=""
    initdict={}
    initdict = ut.get_tokens()
    hf_token = initdict["hf_token"]
    embedding_model_id = initdict["embedding_model"]
    chromadbpath = initdict["dataset_chroma_db"]
    chromadbcollname = initdict["dataset_chroma_db_collection_name"]
    llm_repo_id = initdict["llm_repoid"]
    
    embedding_model = SentenceTransformer(embedding_model_id)
    #print(chromadbpath)
    #print(chromadbcollname)
    chroma_client = chromadb.PersistentClient(path = chromadbpath)
    collection = chroma_client.get_collection(name = chromadbcollname)
 
    #collection = chroma_client.get_or_create_collection(name=chromadbcollname)
    query_vector = embedding_model.encode(query).tolist()
    output = collection.query(
        query_embeddings=[query_vector],
        n_results=1,
        #where={"distances": "is_less_than_1"},
        include=['documents','distances'],
        
        )
    #Filter for distances
    output = filterdistance(output)
    
    if len(output)>0: 
        template = """
        <s>[INST] <<SYS>>
        Act as a patent assistant who is helping summarize and neatly format the results for better readability. Ensure the output is gramatically correct and easily understandable 
        <</SYS>>

        {text} [/INST]
        """
        #Build the prompt template
        prompt = PromptTemplate(
            input_variables=["text"],
            template=template,
        )
        text = output
        
        llm = HuggingFaceHub(huggingfacehub_api_token=hf_token, 
                        repo_id=llm_repo_id, model_kwargs={"temperature":0.2, "max_new_tokens":50})

        result = llm.invoke(prompt.format(text=text))
        print (result)
    return result
  
    return output
    # extract and apply distance condition

st.title("BIG Patent Search")

# Main chat form
with st.form("chat_form"):
    query = st.text_input("Enter the abstract search for similar patents: ")
    #LLM_Summary = st.checkbox('Summarize results with LLM') 
    submit_button = st.form_submit_button("Send")    

if submit_button:
    st.write("Fetching results..\n")
    results =  get_collections(query)
    
    if len(results)>0: 
        #docids = results["documents"]
        response = "There are existing patents related to -    "
        substring = results.partition("[/ASSistant]")[-1]
        if len(substring)>0:
            response  = response + str(substring)
        else: 
            response = response + results.partition("[/INST]")[-1]
            
    else: response = "No results"
    
    st.write (response)