import streamlit as st
import srsly

def search(query):
    results = []
    for grant in grants:
        if query in grant["tags"]:
            results.append({
                "title": grant["title"],
                "tags": grant["tags"]
            })
    st.session_state["results"] = results

st.header("Search 🔎 grants using MeSH 🔖")
st.sidebar.header("Information ℹ")
st.sidebar.write("A complete list of MeSH tags can be found here https://meshb.nlm.nih.gov/treeView")
st.sidebar.write("The grants data can be found https://www.threesixtygiving.org/")
st.sidebar.write("The model used to tag grants is https://huggingface.co/Wellcome/WellcomeBertMesh")

if "grants" not in st.session_state:
    st.session_state["grants"] = list(srsly.read_jsonl("tagged_grants.jsonl"))

grants = st.session_state["grants"]

query = st.text_input("", value="Malaria")
st.button("Search 🔎", on_click=search, kwargs={"query": query})

if "results" in st.session_state:
    st.caption("Related MeSH terms")
    unique_tags = list(set(list([tag for res in st.session_state["results"] for tag in res["tags"]])))
    columns = st.columns(5)
    for row_i in range(3):
        for col_i, col in enumerate(columns):
            with col:
                tag_i = row_i*5 + col_i
                if tag_i < len(unique_tags):
                    tag = unique_tags[tag_i]
                    st.button(tag, on_click=search, kwargs={"query": tag})
    st.table(st.session_state["results"])