#!/usr/bin/env python3
"""
Aspiration Extractor & Clustering Script
----------------------------------------
Reads a CSV with a 'comment' column, extracts user aspirations,
clusters them using embeddings, and saves results to CSV.

Usage:
    python aspiration_extractor.py input.csv [num_clusters]

Dependencies:
    pip install pandas spacy sentence-transformers scikit-learn
    python -m spacy download en_core_web_sm
"""

import sys
import re
import pandas as pd
import spacy
from sentence_transformers import SentenceTransformer
from sklearn.cluster import KMeans

# -------------------------------
# Configuration
# -------------------------------
DEFAULT_CLUSTERS = 8  # fallback if user doesn’t specify
ASPIRATION_LEMMAS = {
    "want", "wish", "hope", "dream", "aspire", "aim", "plan",
    "intend", "desire", "crave", "need", "strive", "determine", "commit"
}

ASPIRATION_PHRASES = [
    r"would like to",
    r"would love to",
    r"my goal is",
    r"looking forward to",
    r"set out to",
    r"i'?m determined to",
    r"i'?m committed to",
    r"going to",
    r"hope to",
    r"plan to",
    r"try to",
    r"want to"
]

# -------------------------------
# Aspiration Extraction
# -------------------------------
def extract_aspiration(text, nlp):
    """Extract the aspiration phrase from a comment text."""
    doc = nlp(text)

    # Lemma-based extraction (verbs like 'want', 'hope')
    for token in doc:
        if token.lemma_.lower() in ASPIRATION_LEMMAS:
            span = doc[token.i + 1 :]  # everything after the verb
            if span.text.strip():
                return span.text.strip()

    # Phrase-based extraction (regex)
    for phrase in ASPIRATION_PHRASES:
        m = re.search(phrase, text.lower())
        if m:
            extracted = text[m.end() :].strip()
            if extracted:
                return extracted

    return None

# -------------------------------
# Main pipeline
# -------------------------------
def main():
    if len(sys.argv) < 2:
        print("Usage: python aspiration_extractor.py input.csv [num_clusters]")
        sys.exit(1)

    input_file = sys.argv[1]
    num_clusters = int(sys.argv[2]) if len(sys.argv) > 2 else DEFAULT_CLUSTERS

    # Load CSV
    try:
        df = pd.read_csv(input_file)
    except Exception as e:
        print(f"❌ Error reading CSV: {e}")
        sys.exit(1)

    if "comment" not in df.columns:
        print("❌ CSV must contain a 'comment' column")
        sys.exit(1)

    comments = df["comment"].dropna().astype(str).tolist()

    print(f"🔹 Loaded {len(comments)} comments from {input_file}")

    # Load models
    print("🔹 Loading NLP models...")
    nlp = spacy.load("en_core_web_sm", disable=["ner", "textcat"])
    embedder = SentenceTransformer("all-MiniLM-L6-v2")

    # Extract aspirations
    extracted = []
    for text in comments:
        aspiration = extract_aspiration(text, nlp)
        extracted.append(aspiration)

    df["extracted_aspiration"] = extracted
    df = df.dropna(subset=["extracted_aspiration"])

    print(f"✅ Found {len(df)} aspiration-related comments")

    if df.empty:
        print("⚠️ No aspirations extracted. Exiting.")
        sys.exit(0)

    # Encode & cluster
    embeddings = embedder.encode(df["extracted_aspiration"].tolist(), show_progress_bar=True)

    n_clusters = min(num_clusters, len(df))  # can't have more clusters than data
    kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
    labels = kmeans.fit_predict(embeddings)

    df["cluster"] = labels

    # Save results
    output_file = "aspiration_clusters.csv"
    df[["comment", "extracted_aspiration", "cluster"]].to_csv(output_file, index=False)

    print(f"💾 Results saved to {output_file}")

    # Preview clusters
    print("\n--- Cluster Preview ---")
    for cluster_id in sorted(df["cluster"].unique()):
        examples = df[df["cluster"] == cluster_id]["extracted_aspiration"].head(3).tolist()
        print(f"\nCluster {cluster_id}:")
        for ex in examples:
            print(f" - {ex}")

# -------------------------------
if __name__ == "__main__":
    main()