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

st.set_page_config(layout="wide")
DATA_FILE = "data/gwf_2017-2021_specter2_base.json"
THEMES = {"cluster": "fall", "year": "mint", "source": "phase"}


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])
    df["year"] = df["year"].replace("", 0)
    df["year"] = df["year"].astype(int)
    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_dataframe():
    return load_df(DATA_FILE)


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

with st.sidebar:
    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)),
    )
    src = st.text_input("Source")

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

    title = st.text_input("Title")

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

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

    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]}")

    color = st.selectbox("Color", ("cluster", "source"))


fig = px.scatter(
    DF,
    x="x",
    y="y",
    opacity=DF["opacity"],
    color=color,
    width=1000,
    height=800,
    hover_data=["title", "authors", "year", "source", "type"],
    color_continuous_scale=THEMES[color],
)
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)