import numpy as np
import streamlit_scrollable_textbox as stx
from ast import literal_eval
import pinecone
import streamlit as st
import json

st.set_page_config(layout="wide")  # isort: split


from utils import nltkmodules
from utils.models import (
    get_bm25_model,
    tokenizer,
    get_data,
    get_instructor_embedding_model,
    get_instructor_embedding_model_api,
    preprocess_text,
)
from utils.retriever import (
    query_pinecone,
    format_context,
    format_query,
    get_bm25_search_hits,
    retrieve_transcript,
)


st.title("Instructor-XL Embeddings")


st.write(
    """The app compares the performance of different instructions using the Instructor-XL Embedding Model on the text from AMD's Q1 2020 Earnings Call Transcript. 
    The app uses a two stage retreival process: 
    1. BM-25 to filter the results based on keyword matching, 
    2. Instructor-XL to perform Semantic Search.
    """
)

use_bm25 = st.checkbox('Use BM25 for filtering results')


data = get_data()


col1, col2 = st.columns([3, 3], gap="medium")

#instructor_model = get_instructor_embedding_model()
instructor_model_api = get_instructor_embedding_model_api()


question_choice = [
    "What was discussed regarding Ryzen revenue performance?",
    "What is the impact of the enterprise and cloud on AMD's growth",
    "What was the impact of situation in China on the sales and revenue?",
]

question_instruction_choice = [
    "Represent the financial question for retrieving supporting documents:",
    "Represent the financial question for retrieving supporting sentences:",
    "Represent the finance query for retrieving supporting documents:",
    "Represent the finance query for retrieving related documents:",
    "Represent a finance query for retrieving relevant documents:",
]


with col1:
    st.subheader("Question")
    st.write(
        "Choose a preset question example from the dropdown or enter a question in the text box."
    )
    default_query = st.selectbox("Question Examples", question_choice)

    query_text = st.text_area(
        "Question",
        value=default_query,
    )

    st.subheader("Question Embedding-Instruction")
    st.write(
        "Choose a preset instruction example from the dropdown or enter a instruction in the text box."
    )
    default_query_embedding_instruction = st.selectbox(
        "Question Embedding-Instruction Examples", question_instruction_choice
    )

    query_embedding_instruction = st.text_area(
        "Question Embedding-Instruction",
        value=default_query_embedding_instruction,
    )

    num_results = int(
        st.number_input("Number of Results to query", 1, 15, value=5)
    )


corpus, bm25 = get_bm25_model(data)

tokenized_query = preprocess_text(query_text).split()
sparse_scores = np.argsort(bm25.get_scores(tokenized_query), axis=0)[::-1]
indices = get_bm25_search_hits(corpus, sparse_scores, 50)


json_output_embedding = instructor_model_api.predict(
    query_embedding_instruction,
    query_text,
    api_name="/predict",
)

json_file = open(json_output_embedding, "r")
json_dict = json.load(json_file)
dense_array = np.array(json_dict["data"], dtype=np.float64)
dense_embedding_api = dense_array.tolist()


#dense_embedding = instructor_model.encode([[query_embedding_instruction, query_text]]).tolist()


text_embedding_instructions_choice = [
    "Represent the financial statement for retrieval:",
    "Represent the financial document for retrieval:",
    "Represent the finance passage for retrieval:",
    "Represent the earnings call transcript for retrieval:",
    "Represent the earnings call transcript sentence for retrieval:",
    "Represent the earnings call transcript answer for retrieval:",
]

index_mapping = {
    "Represent the financial statement for retrieval:": "week14-instructor-xl-amd-fsr-1",
    "Represent the financial document for retrieval:": "week14-instructor-xl-amd-fdr-2",
    "Represent the finance passage for retrieval:": "week14-instructor-xl-amd-fpr-3",
    "Represent the earnings call transcript for retrieval:": "week14-instructor-xl-amd-ectr-4",
    "Represent the earnings call transcript sentence for retrieval:": "week14-instructor-xl-amd-ects-5",
    "Represent the earnings call transcript answer for retrieval:": "week14-instructor-xl-amd-ecta-6",
}

with col2:
    st.write(
        "Choose an instruction for creating the text embedding"
    )
    text_embedding_instruction = st.selectbox(
        "Select instruction for Text Embedding",
        text_embedding_instructions_choice,
    )
    pinecone_index_name = index_mapping[text_embedding_instruction]
    pinecone.init(
        api_key=st.secrets[f"pinecone_{pinecone_index_name}"],
        environment="asia-southeast1-gcp-free",
    )

    pinecone_index = pinecone.Index(pinecone_index_name)
    
    if use_bm25==True:
        matches = query_pinecone(
            dense_vec=dense_embedding_api, top_k=num_results, index=pinecone_index, indices=indices
        )
    else:
        matches = query_pinecone(
            dense_vec=dense_embedding_api, top_k=num_results, index=pinecone_index, indices=None
        )
    context = format_query(matches)
    output_text = format_context(context)

    st.subheader("Retrieved Text:")
    for output in output_text:
        output = f"""{output}"""
        st.write(
            f"<ul><li><p>{output}</p></li></ul>",
            unsafe_allow_html=True,
        )


file_text = retrieve_transcript()
with st.expander("See Transcript"):
    st.subheader("AMD Q1 2020 Earnings Call Transcript:")
    stx.scrollableTextbox(
        file_text, height=700, border=False, fontFamily="Helvetica"
    )