from collections import Counter
import pandas as pd
import streamlit as st
import json
from plotly import express as px
from safetensors import safe_open
from semantic_search import predict
from sentence_transformers import SentenceTransformer
import os
import plotly.graph_objects as go

HF_TOKEN = os.environ.get("HF_TOKEN")


def get_tree_map_data(
    data: dict,
    countings_parents: dict,
    countings_labels: dict,
    root: str = " ",
) -> tuple:
    names: list = [""]
    parents: list = [root]
    values: list = ["0"]

    for group, labels in data.items():
        parents.append(root)
        if group in countings_parents:
            values.append(str(countings_parents[group]))
            group_name_with_count = (
                group
                + "<br>"
                + "Anzahl Datensätze:"
                + " "
                + str(countings_parents[group])
            )
            names.append(group_name_with_count)
        else:
            values.append("0")
            group_name_with_count = group + "<br>" + "Anzahl Datensätze:" + " " + "0"
            names.append(group_name_with_count)
        for label in labels:
            if "-" in label:
                label = label.split("-")
                label = label[0] + "<br> -" + label[1]
            if label in countings_labels:
                label_name_with_count = (
                    label
                    + "<br>"
                    + "<br>"
                    + "Anzahl Datensätze:"
                    + "<br>"
                    + ""
                    + str(countings_labels[label])
                )
                names.append(label_name_with_count)
                parents.append(group_name_with_count)
                values.append(str(countings_labels[label]))
            if label not in countings_labels:
                if "<br>" in label:
                    if (
                        label.split("<br>")[0].strip() + label.split("<br>")[-1]
                        in countings_labels
                    ):
                        label_name_with_count = (
                            label
                            + "<br>"
                            + "<br>"
                            + "Anzahl Datensätze:"
                            + "<br>"
                            + ""
                            + str(
                                countings_labels[
                                    label.split("<br>")[0].strip()
                                    + label.split("<br>")[-1]
                                ]
                            )
                        )
                else:
                    print(label)
                    label_name_with_count = (
                        label
                        + "<br>"
                        + "<br>"
                        + "Anzahl Datensätze:"
                        + "<br>"
                        + ""
                        + "0"
                    )
                    names.append(label_name_with_count)
                    parents.append(group_name_with_count)
                    values.append("0")
    return parents, names, values


def load_json(path: str) -> dict:
    with open(path, "r") as fp:
        return json.load(fp)


# Load Data
data = load_json("data.json")
taxonomy = load_json("taxonomy_processed_v3.json")
taxonomy_labels = [el["group"] + " - " + el["label"] for el in taxonomy]

theme_counts = dict(Counter([el["THEMA"] for el in data]))
labels_counts = dict(Counter([el["BEZEICHNUNG"] for el in data]))

names = [""]
parents = ["Musterdatenkatalog"]

taxonomy_group_label_mapper: dict = {el["group"]: [] for el in taxonomy}

for el in taxonomy:
    if el["group"] != "Sonstiges":
        taxonomy_group_label_mapper[el["group"]].append(el["label"])
    else:
        taxonomy_group_label_mapper[el["group"]].append("Sonstiges ")

del taxonomy_group_label_mapper["Sonstiges"]

parents, names, values = get_tree_map_data(
    data=taxonomy_group_label_mapper,
    countings_parents=theme_counts,
    countings_labels=labels_counts,
    root="Musterdatenkatalog",
)

df = pd.DataFrame(data={"thema": parents, "bezeichnung": names, "value": values})
df["value"] = df["value"].astype(str)
df["bezeichnung"] = df["bezeichnung"]

fig = go.Figure(
    go.Treemap(
        labels=df["bezeichnung"],
        parents=df["thema"],
        textinfo="label",
    )
)
fig.update_layout(margin=dict(t=50, l=25, r=25, b=25))
fig.update_layout(height=1000, width=1000, template="plotly")


tensors = {}
with safe_open("corpus_embeddings.pt", framework="pt", device="cpu") as f:
    for k in f.keys():
        tensors[k] = f.get_tensor(k)

model = SentenceTransformer(
    model_name_or_path="and-effect/musterdatenkatalog_clf",
    device="cpu",
    use_auth_token=HF_TOKEN,
)


st.set_page_config(layout="wide")

st.title("Musterdatenkatalog")

st.markdown(
    """
<style>
.font {
    font-size:20px !important;
}
</style>
""",
    unsafe_allow_html=True,
)

st.markdown(
    '<p class="font">This demo showcases the algorithm of Musterdatenkatalog (MDK) of the Bertelsmann Stiftung. The MDK is a taxonomy of Open Data in municipalities in Germany. It is intended to help municipalities in Germany, as well as data analysts and journalists, to get an overview of the topics and the extent to which cities have already published data sets.</p>',
    unsafe_allow_html=True,
)


st.markdown(
    '<p class="font"> For more details checkout the <a href=https://www.bertelsmann-stiftung.de/de/unsere-projekte/smart-country/musterdatenkatalog> Musterdatenkatalog.</p>',
    unsafe_allow_html=True,
)


col1, col2, col3 = st.columns(3)
col1.metric("Datensätze", len(data))
col2.metric("Themen", len(theme_counts))
col3.metric("Bezeichnungen", len(labels_counts))

st.title("Taxonomy")

st.plotly_chart(fig)

st.title("Predict a Dataset")

st.markdown(
    """
<style>
/* Style columns */
[data-testid="column"] {
      border-radius: 15px;
         background-color: white;
         box-shadow: 0 0 10px #eee;
         border: 1px solid #ddd;
         padding: 1rem;;
} 

/* Style containers */
[data-testid="stVerticalBlock"] > [style*="flex-direction: column;"] > [data-testid="stVerticalBlock"] {
      border-radius: 15px;
         background-color: white;
         box-shadow: 0 0 10px #eee;
         border: 1px solid #ddd;
         padding: 1rem;;
}
</style>
""",
    unsafe_allow_html=True,
)


col1, col2 = st.columns([1.2, 1])


with col2:
    st.subheader("Example Input Dataset Names")
    examples = [
        "Spielplätze",
        "Berliner Weihnachtsmärkte 2022",
        "Hochschulwechslerquoten zum Masterstudium nach Bundesländern",
        "Umringe der Bebauungspläne von Etgert",
    ]

    for example in examples:
        if st.button(example):
            if "key" not in st.session_state:
                st.session_state["query"] = example


with col1:
    if "query" not in st.session_state:
        query = st.text_input(
            "Enter dataset name",
        )
    if "query" in st.session_state and st.session_state.query in examples:
        query = st.text_input("Enter dataset name", value=st.session_state.query)
    if "query" in st.session_state and st.session_state.query not in examples:
        del st.session_state["query"]
        query = st.text_input("Enter dataset name")

    top_k = st.select_slider("Top Results", options=[1, 2, 3, 4, 5], value=1)

    predictions = predict(
        query=query,
        corpus_embeddings=tensors["corpus_embeddings"],
        corpus_labels=taxonomy_labels,
        top_k=top_k,
        model=model,
    )

    if st.button("Predict"):
        for prediction in predictions:
            st.write(prediction)