"""

Run via: streamlit run app.py

"""

import json
import logging

import requests
import streamlit as st
import torch
from datasets import load_dataset
from datasets.dataset_dict import DatasetDict
from transformers import AutoTokenizer, AutoModel

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    level=logging.INFO,
)
logger = logging.getLogger(__name__)

model_hub_url = 'https://huggingface.co/malteos/aspect-scibert-task'

about_page_markdown = f"""# 🔍 Find Papers With Similar Task

See 
- GitHub: https://github.com/malteos/aspect-document-embeddings
- Paper: #TODO
- Model hub: https://huggingface.co/malteos/aspect-scibert-task

"""

# Page setup
st.set_page_config(
    page_title="Papers with similar Task",
    page_icon="🔍",
    layout="centered",
    initial_sidebar_state="auto",
    menu_items={
        'Get help': None,
        'Report a bug': None,
        'About': about_page_markdown,
    }
)

aspects = [
    'task', 'method', 'dataset'
]
tokenizer_name_or_path = f'malteos/aspect-scibert-{aspects[0]}'  # any aspect
dataset_config = 'malteos/aspect-paper-metadata'

tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)


@st.cache(show_spinner=False)
def st_load_model(name_or_path):
    with st.spinner(f'Loading the model `{name_or_path}` (this might take a while)...'):
        model = AutoModel.from_pretrained(name_or_path)
    return model


@st.cache(show_spinner=False)
def st_load_dataset(name_or_path):
    with st.spinner('Loading the dataset (this might take a while)...'):
        dataset = load_dataset(name_or_path)

        if isinstance(dataset, DatasetDict):
            dataset = dataset['train']

        # load existing faiss
        for a in aspects:
            dataset.load_faiss_index(f'{a}_embeddings', f'{a}_embeddings.faiss')

        # add faiss
        #dataset.add_faiss_index(column=f'{aspect}_embeddings')
        #loaded_dataset.add_faiss_index(column='method_embeddings')
        #loaded_dataset.add_faiss_index(column='dataset_embeddings')

    return dataset


aspect_to_model = dict(
    task=st_load_model('malteos/aspect-scibert-task'),
    method=st_load_model('malteos/aspect-scibert-method'),
    dataset=st_load_model('malteos/aspect-scibert-dataset'),
)
dataset = st_load_dataset(dataset_config)


def get_paper(doc_id):
    res = requests.get(f'https://api.semanticscholar.org/v1/paper/{doc_id}')

    if res.status_code == 200:
        return res.json()
    else:
        raise ValueError(f'Cannot load paper from S2 API: {doc_id}')


def find_related_papers(paper_id, user_aspect):
    # Add result to session

    paper = get_paper(paper_id)

    if paper is None or 'title' not in paper or 'abstract' not in paper:
        raise ValueError('Could not retrieve data for input paper')

    title_abs = paper['title'] + ': ' + paper['abstract']

    # preprocess the input
    inputs = tokenizer(title_abs, padding=True, truncation=True, return_tensors="pt", max_length=512)

    # inference
    outputs = aspect_to_model[user_aspect](**inputs)

    # logger.info(f'attention_mask: {inputs["attention_mask"].shape}')
    #
    # logger.info(f'Outputs: {outputs["last_hidden_state"]}')
    # logger.info(f'Outputs: {outputs["last_hidden_state"].shape}')

    # Mean pool the token-level embeddings to get sentence-level embeddings
    embeddings = torch.sum(
        outputs["last_hidden_state"] * inputs['attention_mask'].unsqueeze(-1), dim=1
    ) / torch.clamp(torch.sum(inputs['attention_mask'], dim=1, keepdims=True), min=1e-9)

    result = dict(
        paper=paper,
        aspect=user_aspect,
    )

    result.update(dict(
        #embeddings=embeddings.tolist(),
    ))

    # Retrieval
    prompt = embeddings.detach().numpy()[0]
    scores, retrieved_examples = dataset.get_nearest_examples(f'{user_aspect}_embeddings', prompt, k=10)

    result.update(dict(
        related_papers=retrieved_examples,
    ))

    # st.session_state.results.append(result)

    return result


# # Start session
# if 'results' not in st.session_state:
#     st.session_state.results = []

# Page
st.title('Aspect-based Paper Similarity')
st.markdown("""This demo showcases [Specialized Document Embeddings for Aspect-based Research Paper Similarity](#TODO).""")

# Introduction
st.markdown(f"""The model was trained using a triplet loss on machine learning papers from the [paperswithcode.com](https://paperswithcode.com/) corpus with the objective of pulling embeddings of papers with the same task, method, or datasetclose together. For a more comprehensive overview of the model check out the [model card on 🤗 Model Hub]({model_hub_url}) or read [our paper](#TODO).
""")
st.markdown("""Enter a ArXiv ID or a DOI of a paper for that you want find similar papers.

Try it yourself! 👇""",
    unsafe_allow_html=True)

# Demo
with st.form("aspect-input", clear_on_submit=False):
    paper_id = st.text_input(
        label='Enter paper ID (format "arXiv:<arxiv_id>", "<doi>", or "ACL:<acl_id>"):',
        # value="arXiv:2202.06671",
        placeholder='Any DOI, ACL, or ArXiv ID'
    )

    example = st.selectbox(
        label='Or select example',
        options=[
            "arXiv:2202.06671",
            '10.1016/j.eswa.2019.06.026'
        ]
    )

    # click_clear = st.button('clear text input', key=1)
    # if click_clear:
    #     paper_id = st.text_input(
    #         label='Enter paper ID (arXiv:<arxiv_id>, or <doi>):', value="XXX", placeholder='123')

    user_aspect = st.radio(
        label="In what aspect are you interested?",
        options=aspects
    )

    cols = st.columns(3)
    submitted = cols[1].form_submit_button("Find related papers")

# Listener
if submitted:
    if paper_id or example:
        with st.spinner('Finding related papers...'):
            try:
                result = find_related_papers(paper_id if paper_id else example, user_aspect)

                input_paper = result['paper']
                related_papers = result['related_papers']

                # with st.empty():

                st.markdown(
                    f'''Your input paper: \n\n<a href="{input_paper['url']}"><b>{input_paper['title']}</b></a> ({input_paper['year']})<hr />''',
                    unsafe_allow_html=True)

                related_html = '<ul>'

                for i in range(len(related_papers['paper_id'])):
                    related_html += f'''<li><a href="{related_papers['url_abs'][i]}">{related_papers['title'][i]}</a></li>'''

                related_html += '</ul>'

                st.markdown(f'''Related papers with similar {result['aspect']}: {related_html}''', unsafe_allow_html=True)

            except (TypeError, ValueError, KeyError) as e:
                st.error(f'**Error**: {e}')

    else:
        st.error('**Error**: No paper ID provided. Please provide a ArXiv ID or DOI.')

# # Results
# if 'results' in st.session_state and st.session_state.results:
#     first = True
#     for result in st.session_state.results[::-1]:
#         if not first:
#             st.markdown("---")
#         # st.markdown(f"ID:\n> {result['paperId']}")
#         # col_1, col_2, col_3 = st.columns([1,2,2])
#         # col_1.metric(label='', value=json.dumps(result))
#         # col_2.metric(label='Label', value=f"fooo")
#         # col_3.metric(label='Score', value=f"123")
#         input_paper = result['paper']
#         related_papers = result['related_papers']
#
#         # with st.empty():
#
#         st.markdown(f'''Your input paper: \n\n<a href="{input_paper['url']}"><b>{input_paper['title']}</b></a> ({input_paper['year']})<hr />''', unsafe_allow_html=True)
#
#         related_html = '<ul>'
#
#         for i in range(len(related_papers['paper_id'])):
#             related_html += f'''<li><a href="{related_papers['url_abs'][i]}">{related_papers['title'][i]}</a></li>'''
#
#         related_html += '</ul>'
#
#         st.markdown(f'''Related papers with similar {result['aspect']}: {related_html}''', unsafe_allow_html=True)
#
#         # st.markdown(f'''Related papers: {related_html}''', unsafe_allow_html=True)
#
#         first = False