# %% [markdown]
# # Visualize Most Prepared Users
# This workflow loads user embeddings, generates a prompt embedding, computes similarity, and visualizes the most prepared users.

# %%
import pandas as pd
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import matplotlib.pyplot as plt
import seaborn as sns
from utils import get_ollama_embedding

# %% [markdown]
# ## Read in the file with user vectors
# Update the path/format as needed.
# %%
df = (
    pd.read_csv("./data/eurobarometer_preparedness_model_data_v2.csv").assign(
        **{
            "user_id": lambda x: x["country_code_iso_3166"]
            + "_"
            + x["age_recoded_6_categories"].astype(str)
            + "_"
            + x["gender"].astype(str)
        }
    )
    # select only relevant countries for visualization
    .loc[
        lambda x: x["country_code_iso_3166"].isin(
            ["FI", "DE-E", "DE-W", "FR", "ES", "PT"]
        )
    ]
)
df.shape


# %% [markdown]
# ## Convert string embeddings to lists if needed
# %%
def parse_embedding(x):
    if isinstance(x, str):
        return [float(i) for i in x.strip("[]").split(",")]
    return x


df["embedding"] = df["embedding"].apply(parse_embedding)

# %%
# Select best user:
# disaster_measures_in_hh_* == 1, and how_many_days_meet_* == (4 | 5)
disaster_measures_in_hh_columns = [
    col for col in df.columns if col.startswith("disaster_measures_in_hh_")
]
how_many_days_meet_columns = [
    col for col in df.columns if col.startswith("how_many_days_meet_")
]
best_users = df[
    (df[disaster_measures_in_hh_columns] == 1).all(axis=1)
    & (df[how_many_days_meet_columns].isin([4, 5]).all(axis=1))
]
print(f"Number of best users: {len(best_users)}")

# %% [markdown]
# ## Write your prompt and generate its embedding
# %%
# prompt = "The user is highly prepared for disasters, with emergency supplies and a clear plan."
prompt = best_users.head(1)["user_text"].values[
    0
]  # Example: use the first user's text as the prompt
prompt_embedding = get_ollama_embedding(prompt)

# %% [markdown]
# ## Compute similarity between each user and the prompt
# %%
user_embeddings = np.vstack(df["embedding"].values)
prompt_vec = np.array(prompt_embedding).reshape(1, -1)
similarities = cosine_similarity(user_embeddings, prompt_vec).flatten()
df["similarity"] = similarities

# %% [markdown]
# ## Visualize the users by similarity
# %%
_ = plt.figure(figsize=(5, 7))
_ = sns.violinplot(
    x="similarity", y="country_code_iso_3166", hue="country_code_iso_3166", data=df
)
_ = sns.stripplot(
    x="similarity",
    y="country_code_iso_3166",
    data=df,
    hue="country_code_iso_3166",
    alpha=0.8,
    jitter=True,
    linewidth=0.5,
    edgecolor="white",
)
_ = plt.title("Distribution of User Similarities to Preparedness Prompt")
_ = plt.xlabel("Cosine Similarity")

# %%