import logging
import os
import json
import gradio as gr

from langchain import PromptTemplate, LLMChain
from langchain.chains.question_answering import load_qa_chain
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.chat_models import ChatOpenAI

from prompts import PROMPT_EXTRACT_DATE, PROMPT_FED_ANALYST
from filterminutes import search_with_filter
from examples import FedMinutesSearch

# --------------------------Load the sentence transformer and the vector store--------------------------#
model_name = 'sentence-transformers/all-mpnet-base-v2'
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': False}
embeddings = HuggingFaceEmbeddings(model_name=model_name, model_kwargs=model_kwargs, encode_kwargs=encode_kwargs)
vs = FAISS.load_local("MINUTES_FOMC_HISTORY", embeddings)

# --------------------------Import the prompts------------------#
PROMPT_DATE = PromptTemplate.from_template(PROMPT_EXTRACT_DATE)
PROMPT_ANALYST = PromptTemplate.from_template(PROMPT_FED_ANALYST)


# --------------------------Define the qa chain for answering queries--------------------------#
def load_chains(open_ai_key):
    date_extractor = LLMChain(llm=ChatOpenAI(temperature=0, model_name='gpt-3.5-turbo', openai_api_key=open_ai_key),
                              prompt=PROMPT_DATE)
    fed_chain = load_qa_chain(llm=ChatOpenAI(model_name='gpt-3.5-turbo', temperature=0, openai_api_key=open_ai_key),
                              chain_type='stuff', prompt=PROMPT_ANALYST)
    return date_extractor, fed_chain


def get_chain(query, api_key=os.environ['OPENAI_API_KEY']):
    """
    Detects the date, computes similarity, and answers the query using
    only documents corresponding to the date requested.
    The query is first passed to the date extractor to extract the date
    and then to the qa chain to answer the query.
    Parameters
    ----------
    query : str
        Query to be answered.
    api_key : str
        OpenAI API key.

    Returns
        Answer to the query.
    """
    date_extractor, fed_chain = load_chains(api_key)
    logging.info('Extracting the date in numeric format..')
    date_response = date_extractor.run(query)
    if date_response == 'False':
        logging.info('No date elements found. Running the qa without filtering can output incorrect results.')
        return 'No date elements found. Please include temporal references in your query. ' \
               'If you believe this is wrong, please flag below as appropriate.'
    else:
        filter_date = json.loads(date_response)
        logging.info(f'Date parameters retrieved: {filter_date}')
        logging.info('Running the qa with filtered context..')
        filtered_context = search_with_filter(vs, query, init_k=200, step=300, target_k=7, filter_dict=filter_date)
        logging.info(20 * '-' + 'Metadata for the documents to be used' + 20 * '-')
        for doc in filtered_context:
            logging.info(doc.metadata)
    return fed_chain({'input_documents': filtered_context[:7], 'question': query})['output_text']


if __name__ == '__main__':
    app = gr.Interface(fn=get_chain,
                       inputs=[gr.Textbox(lines=2, placeholder="Enter your query", label='Your query'),
                               gr.Textbox(lines=1, placeholder="Your OpenAI API key here",
                                          label='OpenAI Key (optional, for faster response time)')],
                       description='Here, you can use a ChatGPT-powered retrieval augmented generation system to ask questions'
                                   ' about the [minutes]('
                                   'https://www.federalreserve.gov/monetarypolicy.htm) of the Federal'
                                   ' Open Market Committee meetings from March 1936 to June 2023. The answers are  '
                                   'tuned to focus on economic, '
                                   'cultural, financial, and political developments occurring at a given time.'
                                   ' The model actively looks for the presence of date elements in the query '
                                   'and will stop the execution if cannot find them to minimize the risk of model '
                                   'hallucination. Nevertheless, the usual caveats for applications making use of generative AI apply.'
                                   ' Click the query examples below to see some possible outputs from the model.',
                       article='**Disclaimer**: This app is for demonstration purposes only, and no assurance of uninterrupted'
                       ' functionality can be given at this time. Answers may take some'
                               ' time to complete '
                               'during periods of heavy usage. There is still significant work planned ahead. Please be patient :)',
                       analytics_enabled=True,
                       allow_flagging="manual",
                       flagging_options=["error", "ambiguous", "not true"],
                       outputs=gr.Textbox(lines=1, label='Answer'),
                       title='Search the FED minutes archive',
                       examples=FedMinutesSearch,
                       cache_examples=False
                       )
    app.queue(concurrency_count=2)
    app.launch()