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}')
        if 'year' not in filter_date:
            return 'Please add a specific year to the month.'
        elif (int(filter_date['year']) > 2023 and int(filter_date['month']) > 6) or (int(filter_date['year']) == 1936 and int(filter_date['month']) < 3) or int(filter_date['year']) < 1936:
            return 'Date is in the future or it is before the earliest of the publicly available records.' \
                   ' If you believe this is wrong, please flag below as appropriate.'
        filtered_context = search_with_filter(vs, query, init_k=200, step=500, target_k=7, filter_dict=filter_date)
        logging.info(20 * '-' + 'Metadata for the documents to be used' + 20 * '-')
        if len(filtered_context) == 0:
            return 'There is no information in the minutes for the given date. Please check if there were was an FOMC ' \
                   'meeting  on the given date. If you believe this is wrong, please flag below as appropriate.'
        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 heavy use)')],
                       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", "not useful", "not true"],
                       outputs=gr.Textbox(lines=1, label='Answer'),
                       title='Search the FED minutes archive',
                       examples=FedMinutesSearch,
                       cache_examples=True
                       )
    app.queue(concurrency_count=2)
    app.launch()