# explorer.py
import streamlit as st
import plotly.graph_objects as go
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import pickle


def app():
    # ---------------- Load data ----------------
    path = "./data/"
    with open(path + "coordinates3d_tsne_dict.pkl", "rb") as f:
        coordinates3d_dict = pickle.load(f)
    filtered_vocab_df = pd.read_json(path + "filtered_vocab_df.json")
    filtered_vocab_df.set_index("word", inplace=True)
    with open(path + "vectors_dict_comp.pkl", "rb") as f:
        vectors_dict = pickle.load(f)

    # ---------------- UI controls ----------------
    options = list(vectors_dict.keys())
    subcorpus = st.selectbox(
        "What subcorpus are you interested in:",
        options,
        index=options.index("NOSCEMUS - 1501-1550") if "NOSCEMUS - 1501-1550" in options else 0,
    )
    target = st.text_input("What word are you interested in:", value="scientia")
    topn = st.slider("Select the number of nearest neighbours:", 5, 100, 20)

    # ---------------- Helpers ----------------
    def cosine_similarity_matrix(kv, wordlist):
        vectors = np.array([kv[w] for w in wordlist])
        normalized = vectors / np.linalg.norm(vectors, axis=1, keepdims=True)
        sim = np.dot(normalized, normalized.T)
        return pd.DataFrame(sim, index=wordlist, columns=wordlist)

    def plot_similarity_matrix(sim_df, title="Pairwise Vector Similarity"):
        fig, ax = plt.subplots(figsize=(8, 6))
        cax = ax.matshow(sim_df, cmap="Greens")
        fig.colorbar(cax, label="Cosine Similarity")
        ax.set_title(title, fontsize=14, pad=20)
        ax.set_xlabel("Words", fontsize=12)
        ax.set_ylabel("Words", fontsize=12)
        ticks = np.arange(len(sim_df.columns))
        ax.set_xticks(ticks)
        ax.set_xticklabels(sim_df.columns, rotation=90, fontsize=10)
        ax.set_yticks(ticks)
        ax.set_yticklabels(sim_df.index, fontsize=10)
        ax.grid(False)
        fig.tight_layout()
        return fig

    # ---------------- 3-D plot ----------------
    def plot_gen(subcorpus, target, topn, coordinates3d_dict, filtered_vocab_df, vectors_dict):
        subcorpus_short = subcorpus.replace("NOSCEMUS - ", "")
        xs, ys, zs, words = coordinates3d_dict[subcorpus]

        word_dict = filtered_vocab_df.apply(
            lambda row: "<br>wordcount: "
            + (str(row[subcorpus_short]) if subcorpus_short in filtered_vocab_df.columns else "NA")
            + "<br>translation: "
            + row["transl"],
            axis=1,
        ).to_dict()

        kv = vectors_dict[subcorpus]
        nns_tuples = kv.most_similar(target, topn=topn)
        wordlist = [target] + [t[0] for t in nns_tuples]
        sim_scores = [str(1)] + [str(np.round(t[1], 2)) for t in nns_tuples]

        words = [w.lower() for w in words]
        wordlist = [w.lower() for w in wordlist]
        idx = [words.index(w) for w in wordlist if w in words]

        wordlist_xs, wordlist_ys, wordlist_zs = xs[idx], ys[idx], zs[idx]
        colors = ["darkred"] + ["black"] * len(nns_tuples)
        fontsizes = [18] + [14] * len(nns_tuples)

        hover_text = []
        for word, sim in zip(wordlist, sim_scores):
            if word in word_dict:
                hover_text.append(
                    word + word_dict[word] + f"<br>similarity to target ({target}): {sim}"
                )
            else:
                hover_text.append(
                    word
                    + f"<br>wordcount: NA<br>translation: NA<br>similarity to target ({target}): {sim}"
                )

        # ---- split into markers + text traces (same appearance, more stable) ----
        fig = go.Figure()

        # markers: nearly invisible but keep WebGL stable
        fig.add_trace(
            go.Scatter3d(
                x=wordlist_xs,
                y=wordlist_ys,
                z=wordlist_zs,
                mode="markers",
                marker=dict(size=8, color="purple", opacity=0.01),
                hovertext=hover_text,
                hoverinfo="text",
                name="points",
                showlegend=False,
            )
        )

        # text labels: all visible
        fig.add_trace(
            go.Scatter3d(
                x=wordlist_xs,
                y=wordlist_ys,
                z=wordlist_zs,
                mode="text",
                text=wordlist,
                textposition="middle center",
                textfont=dict(size=fontsizes, color=colors, family="Arial"),
                hovertext=hover_text,
                hoverinfo="text",
                name="labels",
                showlegend=False,
            )
        )

        fig.update_layout(
            title="Embeddings",
            scene=dict(
                xaxis=dict(
                    title="",
                    showgrid=False,
                    showline=False,
                    showticklabels=False,
                    zeroline=False,
                    showbackground=False,
                    linecolor="rgba(0,0,0,0)",
                ),
                yaxis=dict(
                    title="",
                    showgrid=False,
                    showline=False,
                    showticklabels=False,
                    zeroline=False,
                    showbackground=False,
                    linecolor="rgba(0,0,0,0)",
                ),
                zaxis=dict(
                    title="",
                    showgrid=False,
                    showline=False,
                    showticklabels=False,
                    zeroline=False,
                    showbackground=False,
                    linecolor="rgba(0,0,0,0)",
                ),
                bgcolor="rgba(255,255,255,0)",
            ),
            paper_bgcolor="rgba(255,255,255,255)",
            plot_bgcolor="rgba(255,255,255,255)",
            width=980,
            height=700,
            margin=dict(l=0, r=0, b=0, t=0),
            hovermode="closest",
            showlegend=False,
            uniformtext_minsize=12,
            uniformtext_mode="hide",
            uirevision="keep",      # keep camera
            scene_dragmode="orbit",  # orbit by default
        )

        sim_df = cosine_similarity_matrix(kv, wordlist[:35])
        nns_df = pd.DataFrame(nns_tuples, columns=["word", "similarity"]).set_index("word")
        fig_sim = plot_similarity_matrix(sim_df, title="Pairwise Vector Similarity Matrix")
        return fig, nns_df, fig_sim

    # ---------------- Render ----------------
    try:
        fig, nns_df, fig_sim = plot_gen(
            subcorpus, target, topn, coordinates3d_dict, filtered_vocab_df, vectors_dict
        )
        st.markdown(f"### 3D projection of the {topn} words most similar to *{target}*")
        st.plotly_chart(
            fig,
            use_container_width=False,  # fixed size = no resize loop
            key="3d_embeddings_plot",
            config={
                "scrollZoom": True,
                "displayModeBar": True,
                "modeBarButtonsToAdd": ["resetCameraDefault3d"],
            },
        )
        st.markdown(f"### {topn} nearest neighbours of *{target}*")
        st.dataframe(nns_df, width=300)
        st.pyplot(fig_sim)
    except Exception as e:
        st.write(f"The word **{target}** is not in our corpus! ({e})")