# -----------------------
# 1. Setup
# -----------------------
import pandas as pd
from allennlp.predictors.predictor import Predictor
import allennlp_models.tagging
from sentence_transformers import SentenceTransformer
import hdbscan
from sklearn.decomposition import PCA
import matplotlib.pyplot as plt
import sys
import os

# ✅ Transformer-only SRL model
predictor = Predictor.from_path(
    "https://storage.googleapis.com/allennlp-public-models/structured-prediction-srl-bert.2020.12.15.tar.gz"
)

# Embedding model
embedder = SentenceTransformer("all-MiniLM-L6-v2")

# -----------------------
# 2. Config
# -----------------------
OUTPUT_FILE = "reddit_jtbd_partial.csv"
FINAL_FILE = "reddit_jtbd_hdbscan.csv"
CHUNK_SIZE = 100  # save progress every N comments

# -----------------------
# 3. Load comments
# -----------------------
df = pd.read_csv(sys.argv[1])
comments = df["comment"].dropna().tolist()

# -----------------------
# 4. Resume logic
# -----------------------
if os.path.exists(OUTPUT_FILE):
    processed_df = pd.read_csv(OUTPUT_FILE)
    processed_indices = set(processed_df["index"].tolist())
    print(f"🔄 Resuming from {len(processed_indices)} already processed rows.")
else:
    processed_df = pd.DataFrame(columns=["index", "comment", "situation", "struggle", "outcome", "job_story"])
    processed_indices = set()

results = []

# -----------------------
# 5. JTBD extraction
# -----------------------
def extract_jtbd_triplet(text):
    if not isinstance(text, str) or not text.strip():
        return None

    try:
        results = predictor.predict(sentence=text)
    except Exception as e:
        print(f"SRL failed on: {text[:50]}... ({e})")
        return None
    
    triplets = []
    
    for verb in results["verbs"]:
        desc = verb["description"]

        situation, struggle, outcome = [], [], []

        spans = desc.split("]")
        for span in spans:
            if "[" not in span or ":" not in span:
                continue  # skip malformed
            try:
                role, phrase = span.split(":", 1)
            except ValueError:
                continue

            role = role.strip("[ ")
            phrase = phrase.strip()

            if role in ["ARGM-TMP", "ARGM-CAU"]:
                situation.append(phrase)
            elif role in ["ARG1", "ARGM-MNR"]:
                struggle.append(phrase)
            elif role in ["ARGM-PRP", "ARGM-PNC"]:
                outcome.append(phrase)

        if situation or struggle or outcome:
            job_story = {
                "situation": " ".join(situation) if situation else None,
                "struggle": " ".join(struggle) if struggle else None,
                "outcome": " ".join(outcome) if outcome else None,
            }
            triplets.append(job_story)

            print(f"Found JTBD → Situation: {job_story['situation']} | "
                  f"Struggle: {job_story['struggle']} | "
                  f"Outcome: {job_story['outcome']}")

    return triplets if triplets else None


# -----------------------
# 6. Process comments
# -----------------------
for i, comment in enumerate(comments):
    if i in processed_indices:
        continue  # skip already done

    triplets = extract_jtbd_triplet(comment)
    if triplets:
        for t in triplets:
            job_story = f"When {t['situation']}, I {t['struggle']}, so I can {t['outcome']}".replace("None", "").strip()
            results.append({
                "index": i,
                "comment": comment,
                "situation": t['situation'],
                "struggle": t['struggle'],
                "outcome": t['outcome'],
                "job_story": job_story
            })

    # Save every CHUNK_SIZE iterations
    if i % CHUNK_SIZE == 0 and results:
        temp_df = pd.DataFrame(results)
        processed_df = pd.concat([processed_df, temp_df], ignore_index=True)
        processed_df.to_csv(OUTPUT_FILE, index=False)
        results = []
        print(f"💾 Saved progress up to row {i}")

# Final save (remaining rows)
if results:
    temp_df = pd.DataFrame(results)
    processed_df = pd.concat([processed_df, temp_df], ignore_index=True)
    processed_df.to_csv(OUTPUT_FILE, index=False)

print("✅ Finished extraction, now embedding + clustering...")

# -----------------------
# 7. Embedding + HDBSCAN clustering
# -----------------------
texts = processed_df["job_story"].dropna().tolist()
embeddings = embedder.encode(texts, show_progress_bar=True)

# HDBSCAN clustering
clusterer = hdbscan.HDBSCAN(min_cluster_size=5, metric='euclidean')
processed_df["cluster"] = clusterer.fit_predict(embeddings)

# -----------------------
# 8. Review clusters
# -----------------------
cluster_labels = sorted(processed_df["cluster"].unique())
for c in cluster_labels:
    if c == -1:
        print("\n🌀 Noise / uncategorized examples:")
    else:
        print(f"\n📌 Cluster {c} theme examples:")
    examples = processed_df[processed_df["cluster"] == c]["job_story"].head(5).tolist()
    for e in examples:
        print(" -", e)

# -----------------------
# 9. Visualization
# -----------------------
pca = PCA(n_components=2)
reduced = pca.fit_transform(embeddings)
plt.figure(figsize=(10, 6))
scatter = plt.scatter(reduced[:,0], reduced[:,1], c=processed_df["cluster"], cmap="tab10", alpha=0.7)
plt.legend(*scatter.legend_elements(), title="Cluster")
plt.title("JTBD Clusters with HDBSCAN (PCA projection)")
plt.show()

# -----------------------
# 10. Save results
# -----------------------
processed_df.to_csv(FINAL_FILE, index=False)
print(f"✅ Saved clustered job stories to {FINAL_FILE}")