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.prompts import PromptTemplate
from langchain.chains.question_answering import load_qa_chain
from langchain.chains.qa_with_sources import load_qa_with_sources_chain
from langchain import LLMMathChain, SQLDatabase, SQLDatabaseChain, LLMChain
from langchain.agents import initialize_agent, Tool

import sqlite3
#import pandas as pd
import json

# database
cxn = sqlite3.connect('./data/mbr.db')

"""# 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

"""# LLM based qa functions"""

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 the user 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_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

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

# for x in xmpl_list:
#   print(32*'=')
#   print(x)
#   print(32*'=')
#   r = get_llm_response(x)
#   print(r)

"""# Database query"""

db = SQLDatabase.from_uri("sqlite:///./data/mbr.db")

llm = OpenAI(temperature=0)
# default model
# model_name: str = "text-davinci-003"
# instruction fine-tuned, sometimes referred to as GPT-3.5

db_chain = SQLDatabaseChain(llm=llm, database=db, verbose=True)

def db_qry(qry):
  responses = db_chain.run(query='my mbr_id is 456 ;'+str(qry) ) ############### hardcode mbr id 456 for demo
  return responses

#db_qry('how many footcare visits have I had?')

"""## Math
- default version
"""

llm_math_chain = LLMMathChain(llm=llm, verbose=True)

#llm_math_chain.run('what is the square root of 49?')

"""# Greeting"""

template = """You are a friendly AI assistant for the insurance company Humana. 
Your name is Bruce and you were created on February 13, 20203.
Offer polite greetings and brief small talk.
Respond to thanks with, 'Glad to help.'
If the question is not about Humana, politely guide the user to ask questions about Humana insurance benefits.
QUESTION: {question}
=========
FINAL ANSWER:"""
greet_prompt = PromptTemplate(template=template, input_variables=["question"])

greet_llm = LLMChain(prompt=greet_prompt, llm=llm, verbose=True)

"""# MRKL Chain"""

tools = [
    Tool(
        name = "Benefit",
        func=get_llm_response,
        description="useful for when you need to answer questions about plan benefits, premiums and payments. You should ask targeted questions"
    ),
    Tool(
        name="Calculator",
        func=llm_math_chain.run,
        description="useful for when you need to answer questions about math"
    ),
    Tool(
        name="Member DB",
        func=db_qry,
        description="useful for when you need to answer questions about member details such their accumulated use of services. Input should be in the form of a question containing full context"
    ),
    Tool(
        name="Greeting",
        func=greet_llm.run,
        description="useful for when you need to respond to greetings, thanks and make small talk"
    ),
]

mrkl = initialize_agent(tools, llm, agent="zero-shot-react-description", verbose=True, return_intermediate_steps=True, max_iterations=5, early_stopping_method="generate")

def mrkl_rspnd(qry):
  response = mrkl({"input":str(qry) })
  return response

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