from collections import Counter
import pandas as pd
import streamlit as st
import json
from safetensors import safe_open
from semantic_search import predict
from sentence_transformers import SentenceTransformer
import os
import plotly.graph_objects as go
import plotly.express as px
from utils.process_data import load_data, merge_geoemtry, add_coor

HF_TOKEN = os.environ.get("HF_TOKEN")
CITIES_ENRICHED = os.path.join("data", "cities_enriched_manually.csv")
DATA = os.path.join("2024-08-21_musterdatenkatalog.json")

TAXONOMY = os.path.join("taxonomy_processed_v3.json")
MAP_PATH = os.path.join("data", "map_data.csv")  # this is only for saving
MAP_PATH_WITH_COORD = os.path.join(
    "data", "map_data_with_coord.csv"
)  # this is for saving the data with coordinates
# and local testing


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)
taxonomy = load_json(TAXONOMY)
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")

# # load data ready to plot for local testing
# germany = pd.read_csv(MAP_PATH)
# germany.drop(columns=["lat", "lon"], inplace=True)
# # or generate it directly in this script
map_data = load_data()
map_data = merge_geoemtry(map_data, pd.read_csv(filepath_or_buffer=CITIES_ENRICHED))
# print(
#     map_data["Geometry"].iloc[0].strip("[]").split(),
#     type(map_data["Geometry"].iloc[0].strip("[]").split()),
# )
germany = add_coor(map_data)
germany.to_csv(MAP_PATH_WITH_COORD, index=False)

# # germany need columns with lat and lon as well as hover data
fig_map = px.scatter_mapbox(
    germany,
    lat="lat",
    lon="lon",
    hover_name="ORG",
    custom_data=["Count"],
    # color_discrete_map=["magenta"],
    zoom=5,
    height=700,
)
# Custom hover template
fig_map.update_traces(
    hovertemplate="<br>".join(
        [
            "Kommune: %{hovertext}",  # Use hover_name as hovertext
            "Count: %{customdata[0]}",  # Access elements in custom_data
        ]
    )
)
fig_map.update_layout(mapbox_style="carto-positron")

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 (MDK)")

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

st.markdown(
    """
<style>
.prediction {
    font-size:10px !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 </a>.</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.header("Explore the MDK-Classifier")

st.markdown(
    '<p class="font"> This section allows you to predict a label from the MDK Taxonomy for a title of a dataset from municipalities. You can either enter your own dataset title or click on one of the examples. Checkout also <a href=https://www.govdata.de/> GOVDATA </a> for more dataset title examples. \
    \
    If you click on predict, the model will predict the most likely label for the dataset title. You can also change the number of labels that should be predicted. For example, if you change the Top Results to 3, the model will predict the 3 most likely labels for the dataset title in descending order. </p>',
    unsafe_allow_html=True,
)

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

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


with col2:
    st.markdown(
        '<p class="example">Example Titles of Datasets</p>',
        unsafe_allow_html=True,
    )
    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:
    tabs_font_css = """
    <style>
    div[class*="stTextInput"] label p {
    font-size: 2px;
    }
    </style>    
    """

    st.write(tabs_font_css, unsafe_allow_html=True)

    st.markdown(
        '<p class="example">Enter a dataset title</p>',
        unsafe_allow_html=True,
    )

    if "query" not in st.session_state:
        query = st.text_input("")
    if "query" in st.session_state and st.session_state.query in examples:
        query = st.text_input("Enter a dataset title", 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 a dataset title")

    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.markdown(f'<p class="font"> {prediction} <p>', unsafe_allow_html=True)


st.header("Musterdatenkatalog Taxonomy")
st.write("Data as of 20.08.2024")
st.plotly_chart(fig)

st.header("Locations with Musterdatensätzen")
st.write("Data as of 20.08.2024")
st.markdown(
    """<p class="font">Hover over the map to see how many datasets are available
            for this location. </p>
   """,
    unsafe_allow_html=True,
)
st.plotly_chart(fig_map)