#!/usr/bin/env python3
import sys
import pandas as pd
import spacy
from sentence_transformers import SentenceTransformer
import umap
import hdbscan
from collections import defaultdict

# --- Config ---
COMMENT_COLUMN = "comment"   # column name in CSV

def main(csv_file):
    # Load data
    df = pd.read_csv(csv_file)
    if COMMENT_COLUMN not in df.columns:
        raise ValueError(f"CSV must contain a '{COMMENT_COLUMN}' column")

    comments = df[COMMENT_COLUMN].dropna().tolist()

    # Load spaCy + SentenceTransformer
    nlp = spacy.load("en_core_web_sm")
    embedder = SentenceTransformer("all-MiniLM-L6-v2")

    # Extract nouns + verbs
    all_words = []
    word_to_comments = defaultdict(list)

    for text in comments:
        doc = nlp(text)
        nouns = [t.lemma_.lower() for t in doc if t.pos_ == "NOUN"]
        verbs = [t.lemma_.lower() for t in doc if t.pos_ == "VERB"]

        for w in nouns + verbs:
            all_words.append(w)
            word_to_comments[w].append(text)

    unique_words = list(set(all_words))
    if not unique_words:
        print("No nouns or verbs found.")
        return

    # Embed words
    embeddings = embedder.encode(unique_words, show_progress_bar=True)

    # Reduce dimensions
    reducer = umap.UMAP(n_neighbors=5, min_dist=0.3, metric="cosine", random_state=42)
    umap_embeddings = reducer.fit_transform(embeddings)

    # Cluster
    clusterer = hdbscan.HDBSCAN(min_cluster_size=3, metric="euclidean")
    labels = clusterer.fit_predict(umap_embeddings)

    # Collect clusters
    clusters = defaultdict(list)
    word_to_cluster = {}
    for word, label in zip(unique_words, labels):
        if label == -1:  # noise
            continue
        clusters[label].append(word)
        word_to_cluster[word] = label

    # --- Save clusters.csv ---
    cluster_rows = []
    for cid, words in clusters.items():
        for w in words:
            cluster_rows.append({"cluster_id": cid, "word": w})
    pd.DataFrame(cluster_rows).to_csv("clusters.csv", index=False)

    print("Saved cluster definitions to clusters.csv")

    # --- Save comment_clusters.csv ---
    comment_cluster_map = []
    for text in comments:
        doc = nlp(text)
        words = [t.lemma_.lower() for t in doc if t.pos_ in {"NOUN", "VERB"}]
        cluster_ids = list({word_to_cluster[w] for w in words if w in word_to_cluster})
        comment_cluster_map.append({
            "comment": text,
            "clusters": ";".join(map(str, cluster_ids)) if cluster_ids else ""
        })

    pd.DataFrame(comment_cluster_map).to_csv("comment_clusters.csv", index=False)

    print("Saved comment → cluster mapping to comment_clusters.csv")

    # --- Print summary ---
    print("\n=== Word Clusters (contexts + behaviors) ===")
    for cid, words in clusters.items():
        print(f"\nCluster {cid}: {words}")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python pos_cluster.py <csv_file>")
        sys.exit(1)
    main(sys.argv[1])