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

OPTS = {
    "ACL 2022": "data/acl2022_spectre2-base.json",
    "ACL 2023": "data/acl2023_spectre2-base.json",
}


def load_df(data_file: os.PathLike):
    df = pd.read_json(data_file, orient="records")
    df["x"] = df["point2d"].apply(lambda x: x[0])
    df["y"] = df["point2d"].apply(lambda x: x[1])
    if "publication_type" in df.columns:
        df["type"] = df["publication_type"]
        df = df.drop(columns=["point2d", "publication_type"])
    else:
        df = df.drop(columns=["point2d"])
    return df


@st.cache_data
def load_dataframes():
    return {venue: load_df(df_file) for venue, df_file in OPTS.items()}

DFS = load_dataframes()

option = st.selectbox(
    'Please select a venue',
    ('ACL 2022', 'ACL 2023',)
)


external_stylesheets = ["https://codepen.io/chriddyp/pen/bWLwgP.css"]


fig = px.scatter(
    DFS[option],
    x="x",
    y="y",
    color="cluster",
    width=1200,
    height=750,
    hover_data=["title", "authors", "year", "source", "type"],
    color_continuous_scale="fall",
)
fig.update_layout(
    # margin=dict(l=10, r=10, t=10, b=10),
    showlegend=False,
    font=dict(
        family="Times New Roman",
        size=30,
    ),
)
fig.update_xaxes(title="")
fig.update_yaxes(title="")

st.plotly_chart(fig, use_container_width=True)