"""

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: https://arxiv.org/abs/2203.14541
- 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,
    }
)

aspect_labels = {
    'task': 'Task 🎯 ',
    'method': 'Method 🔨 ',
    'dataset': 'Dataset 🏷️',
}
aspects = list(aspect_labels.keys())
tokenizer_name_or_path = f'malteos/aspect-scibert-{aspects[0]}'  # any aspect
dataset_config = 'malteos/aspect-paper-metadata'


@st.cache(show_spinner=True)
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=True)
def st_load_dataset(name_or_path):
    with st.spinner('Loading the dataset and search index (this might take a while)...'):
        dataset = load_dataset(name_or_path)

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

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

    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)


@st.cache(show_spinner=True)
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 get_embedding(input_text, user_aspect):
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)

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

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

    # 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)

    return embeddings.detach().numpy()[0]


#@st.cache(show_spinner=False)
def find_related_papers(paper_id, user_aspect):
    with st.spinner('Searching for related papers...'):
        paper_id = paper_id.strip()  # remove white spaces

        paper = get_paper(paper_id)

        if paper is None or 'title' not in paper or paper['title'] is None or 'abstract' not in paper or paper['abstract'] is None:
            raise ValueError(f'Could not retrieve title and abstract for input paper (the paper is probably behind a paywall): {paper_id}')

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

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

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

        # Retrieval
        prompt = get_embedding(title_abs, user_aspect)
        scores, retrieved_examples = dataset.get_nearest_examples(f'{user_aspect}_embeddings', prompt, k=10)

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

    return result


# Page
st.title('Aspect-based Paper Similarity')
st.markdown("""This demo showcases [Specialized Document Embeddings for Aspect-based Research Paper Similarity](https://arxiv.org/abs/2203.14541).""")

# 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 dataset close together. 
For a more comprehensive overview of the model check out the [model card on 🤗 Model Hub]({model_hub_url}) or read [our paper](https://arxiv.org/abs/2203.14541).""")
st.markdown("""Enter a ArXiv ID or a DOI of a paper for that you want find similar papers. The title and abstract of the input paper must be available through the [Semantic Scholar API](https://www.semanticscholar.org/product/api).

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_labels = {
        "arXiv:1902.06818": "Data augmentation for low resource sentiment analysis using generative adversarial networks",
        "arXiv:2202.06671": "Neighborhood Contrastive Learning for Scientific Document Representations with Citation Embeddings",
        "ACL:N19-1423": "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding",
        "10.18653/v1/S16-1001": "SemEval-2016 Task 4: Sentiment Analysis in Twitter",
        "10.1145/3065386": "ImageNet classification with deep convolutional neural networks",
        "arXiv:2101.08700": "Multi-sense embeddings through a word sense disambiguation process",
        "10.1145/3340531.3411878": "Incremental and parallel computation of structural graph summaries for evolving graphs",
    }

    example = st.selectbox(
        label='Or select an example:',
        options=list(example_labels.keys()),
        format_func=lambda option_key: f'{example_labels[option_key]} ({option_key})',
    )

    user_aspect = st.radio(
        label="In what aspect are you interested?",
        options=aspects,
        format_func=lambda option_key: aspect_labels[option_key],
    )

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

# Listener
if submitted:
    if paper_id or example:
        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.')