import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
import numpy as np

# ---------- SETTINGS ----------
INPUT_FILE = "daily_habits_daily_steps_pain_points_painpoint_clusters.csv"
OUTPUT_FILE = "habit_subthemes_keywords.csv"
TOP_N_KEYWORDS = 4   # number of keywords to extract as subtheme

# ---------- LOAD DATA ----------
df = pd.read_csv(INPUT_FILE)
df = df.dropna(subset=["quotes"])

# ---------- EXTRACT KEYWORDS PER CLUSTER ----------
subtheme_map = {}

for cluster_id in df["cluster"].unique():
    cluster_quotes = df[df["cluster"] == cluster_id]["quotes"].tolist()
    
    if len(cluster_quotes) == 0:
        continue
    
    # TF-IDF on quotes in this cluster
    vectorizer = TfidfVectorizer(stop_words="english")
    X = vectorizer.fit_transform(cluster_quotes)
    
    terms = vectorizer.get_feature_names_out()
    sums = np.array(X.sum(axis=0)).flatten()
    
    # Pick top N keywords
    top_ids = sums.argsort()[::-1][:TOP_N_KEYWORDS]
    keywords = [terms[i] for i in top_ids]
    
    subtheme_map[cluster_id] = ", ".join(keywords)

# ---------- ASSIGN SUBTHEME ----------
df["subtheme"] = df["cluster"].map(subtheme_map)

# ---------- SAVE RESULTS ----------
df_out = df[["quotes", "cluster", "subtheme"]]
df_out.to_csv(OUTPUT_FILE, index=False)

print(f"✅ Done! Results saved to {OUTPUT_FILE}")