import time
import json
from pyserini.search.lucene import LuceneImpactSearcher
import streamlit as st
from pathlib import Path
import sys
path_root = Path("./")
sys.path.append(str(path_root))


encoder_index_map = {
    'uniCOIL': ('UniCoil', 'index-unicoil'),
    'SPLADE++ Ensemble Distil': ('SpladePlusPlusEnsembleDistil', 'index-splade-pp-ed'),
    'SPLADE++ Self Distil': ('SpladePlusPlusSelfDistil', 'index-splade-pp-sd')
}

index = 'index-splade-pp-ed'
encoder = 'SpladePlusPlusEnsembleDistil'

st.set_page_config(page_title="Pyserini with ONNX Runtime",
                   page_icon='🌸', layout="centered")

cola, colb, colc = st.columns([5, 4, 5])
with colb:
    st.image("logo.jpeg")

colaa, colbb, colcc = st.columns([1, 8, 1])
with colbb:
    encoder = st.select_slider(
        'Select a query encoder with ONNX Runtime',
        options=['uniCOIL', 'SPLADE++ Ensemble Distil', 'SPLADE++ Self Distil'])
    st.write('Now Running Encoder: ', encoder)

encoder, index = encoder_index_map[encoder]

col1, col2 = st.columns([9, 1])
with col1:
    search_query = st.text_input(label="search query", placeholder="Search")

with col2:
    st.write('#')
    button_clicked = st.button("🔎")

searcher = LuceneImpactSearcher(
    f'indexes/{index}', f'{encoder}', encoder_type='onnx')

if search_query or button_clicked:
    num_results = None
    t_0 = time.time()
    print("search query is:\t", search_query)
    search_results = searcher.search(search_query, k=10)
    search_time = time.time() - t_0
    st.write(
        f'<p align=\"right\" style=\"color:grey;\">Retrieved {len(search_results):,.0f} documents in {search_time*1000:.2f} ms</p>', unsafe_allow_html=True)
    for i, result in enumerate(search_results[:10]):
        result_score = result.score
        result_id = result.docid
        output = f'<div class="row"> <div class="column"> <b>Rank</b>: {i+1} </div><div class="column"><b>Document ID</b>: {result_id}</div><div class="column"><b>Score</b>:{result_score:.2f}</div></div>'

        try:
            st.write(output, unsafe_allow_html=True)

        except:
            pass
        st.write('---')