"""
Reference: https://huggingface.co/spaces/gwf-uwaterloo/acl-spectrum (By Ehsan Khamallo)
"""

import os
import re
import pandas as pd
import plotly.express as px
import streamlit as st

st.set_page_config(layout="wide")
DATA_FILE = "hess_papers_details.json"

st.markdown(
    """
    <link href="https://cdn.jsdelivr.net/npm/bootstrap@4.6.1/dist/css/bootstrap.min.css" rel="stylesheet" integrity="sha256-DF7Zhf293AJxJNTmh5zhoYYIMs2oXitRfBjY+9L//AY=" crossorigin="anonymous">
    <link rel="preconnect" href="https://fonts.googleapis.com">
    <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
    <link href="https://fonts.googleapis.com/css2?family=Permanent+Marker&display=swap" rel="stylesheet">
    <style>
    .title {
        font-family: 'Arial';
        font-size: 2.0rem;
    }
    </style>""",
    unsafe_allow_html=True,
)

st.sidebar.write(
    """<center><p class="title">
    Clustering on HESS Papers 🌎🌿
    </p></center>""",
    unsafe_allow_html=True,
)

st.sidebar.write(
    """<p class="text-justify">
    A clustered visualization of all papers submitted to the
    <a href=https://www.hydrology-and-earth-system-sciences.net/>Hydrology and Earth System Sciences</a> (HESS) conference.
    5318 papers are embedded using <a href="https://huggingface.co/allenai/specter2_base">spectre2</a> and reduced with
    t-SNE. Papers span from as early as 1997 to 2023.
    </p>""",
    unsafe_allow_html=True,
)

def to_string_authors(list_of_authors):
    if len(list_of_authors) > 5:
        return ", ".join(list_of_authors[:5]) + ", et al."
    elif len(list_of_authors) > 2:
        return ", ".join(list_of_authors[:-1]) + ", and " + list_of_authors[-1]
    else:
        return " and ".join(list_of_authors)


def load_df(data_file: os.PathLike):
    df = pd.read_json(data_file, orient="records")
    df["x"] = df["t-SNE1"]
    df["y"] = df["t-SNE2"]

    df["authors_trimmed"] = df["authors_trimmed"]

    # #sort dataframe by year
    # df['year'] = pd.to_datetime(df['year'])
    # df = df.sort_values('year', ascending=True)
    # df['year'] = df['year'].dt.strftime('%Y')
    #df['year'] = df['year'].astype(int)

    return df

@st.cache_data
def load_dataframe():
    return load_df(DATA_FILE)

DF = load_dataframe()
DF["opacity"] = 0.04
min_year, max_year = DF["year"].min(), DF["year"].max()

with st.sidebar:
    author_names = st.text_input("Author names (separated by comma)")

    title = st.text_input("Title")

    # Work on this
    # topics = st.multiselect(
    #     "Topics",
    #     ["Topics 1: "],
    #     ["Topics 2: "],
    # )

    start_year, end_year = st.select_slider(
        "Publication year",
        options=[str(y) for y in range(min_year, max_year + 1)],
        value=(str(min_year), str(max_year)),
    )

    start_year = int(start_year)
    end_year = int(end_year)
    df_mask = (DF["year"] >= start_year) & (DF["year"] <= end_year)

    if author_names:
        authors = [a.strip() for a in author_names.split(",")]
        author_mask = DF.authors.apply(
            lambda row: all(any(re.match(rf".*{a}.*", x, re.IGNORECASE) for x in row) for a in authors)
        )
        df_mask = df_mask & author_mask

    if title:
        df_mask = df_mask & DF.title.apply(lambda x: title.lower() in x.lower())

    DF.loc[df_mask, "opacity"] = 1.0
    st.write(f"Number of points: {DF[df_mask].shape[0]}")

fig = px.scatter(
    DF,
    x="x",
    y="y",
    opacity=DF["opacity"],
    color=DF["cluster"],
    width=1000,
    height=800,
    custom_data=("title", "authors_trimmed", "year"),
    color_continuous_scale="haline",
)

fig.update_traces(
hovertemplate="<b>%{customdata[0]}</b><br>%{customdata[1]}<br>%{customdata[2]}<br><i>"
    )

fig.update_layout(
    showlegend=False,
    font=dict(
        family="Times New Roman",
        size=30,
    ),
    hoverlabel=dict(
        align="left",
        font_size=14,
        font_family="Rockwell",
        namelength=-1,
    ),
)

fig.update_xaxes(title="")
fig.update_yaxes(title="")

st.plotly_chart(fig, use_container_width=True)