from sentence_transformers import SentenceTransformer, CrossEncoder, util
from torch import tensor as torch_tensor
from datasets import load_dataset

from langchain.llms import OpenAI
from langchain.docstore.document import Document
from langchain.chains.question_answering import load_qa_chain
from langchain.chains.qa_with_sources import load_qa_with_sources_chain
from langchain.prompts import PromptTemplate


"""# import models"""

bi_encoder = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1')
bi_encoder.max_seq_length = 256     #Truncate long passages to 256 tokens

#The bi-encoder will retrieve top_k documents. We use a cross-encoder, to re-rank the results list to improve the quality
cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')



"""# import datasets"""

dataset = load_dataset("gfhayworth/hack_policy", split='train')
mypassages = list(dataset.to_pandas()['psg'])

dataset_embed = load_dataset("gfhayworth/hack_policy_embed", split='train')
dataset_embed_pd = dataset_embed.to_pandas()
mycorpus_embeddings = torch_tensor(dataset_embed_pd.values)

def search(query, passages = mypassages, doc_embedding = mycorpus_embeddings, top_k=20, top_n = 1):
    question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
    question_embedding = question_embedding #.cuda()
    hits = util.semantic_search(question_embedding, doc_embedding, top_k=top_k)
    hits = hits[0]  # Get the hits for the first query

    ##### Re-Ranking #####
    cross_inp = [[query, passages[hit['corpus_id']]] for hit in hits]
    cross_scores = cross_encoder.predict(cross_inp)

    # Sort results by the cross-encoder scores
    for idx in range(len(cross_scores)):
        hits[idx]['cross-score'] = cross_scores[idx]

    hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
    predictions = hits[:top_n]
    return predictions
    # for hit in hits[0:3]:
    #         print("\t{:.3f}\t{}".format(hit['cross-score'], mypassages[hit['corpus_id']].replace("\n", " ")))



def get_text_fmt(qry, passages = mypassages, doc_embedding=mycorpus_embeddings):
    predictions = search(qry, passages = passages, doc_embedding = doc_embedding, top_n=5, )
    prediction_text = []
    for hit in predictions:
        page_content = passages[hit['corpus_id']]
        metadata = {"source": hit['corpus_id']}
        result = Document(page_content=page_content, metadata=metadata)
        prediction_text.append(result)
    return prediction_text

template = """You are a friendly AI assistant for the insurance company Humana. Given the following extracted parts of a long document and a question, create a succinct final answer. 
If you don't know the answer, just say that you don't know. Don't try to make up an answer.
If the question is not about Humana, politely inform them that you are tuned to only answer questions about Humana.
QUESTION: {question}
=========
{context}
=========
FINAL ANSWER:"""
PROMPT = PromptTemplate(template=template, input_variables=["context", "question"])

chain_qa = load_qa_chain(OpenAI(temperature=0), chain_type="stuff", prompt=PROMPT)


def get_llm_response(message):
    mydocs = get_text_fmt(message)
    response = chain_qa.run(input_documents=mydocs, question=message)
    return response

def chat(message, history):
    history = history or []
    message = message.lower()
    
    response = get_llm_response(message)
    history.append((message, response))
    return history, history