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

# --------------------------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.')
        raise gr.Error('No date elements found. Please include temporal references in in your query.')
    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 ask questions about information in 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 user-defined times.'
                                   'The model actively looks for the presence of date elements present in the query '
                                   'and will raise an error if it cannot find any.'
                                   ' Click the examples below to see the tool in action.',
                       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. Please be patient :)',
                       analytics_enabled=True,
                       outputs=gr.Textbox(lines=1, label='Answer'),
                       title='Query the FED-FOMC meeting minutes',
                       examples=[['Who were the voting members present in the meeting on March 2023?'],
                                 ['What was the economic outlook from the staff presented in the meeting '
                                  'of April 2009 with respect to labour market developments and industrial production?'],
                                 ['How important was the pandemic of Covid-19 in the discussions during 2020 for the effects of monetary policy?'],
                                 ['What was the impact of the oil crisis for the economic outlook during 1973?']],
                       cache_examples=False
                       )
    app.queue()
    app.launch()