#!/usr/bin/env python3
"""
JTBD Keyphrase Clustering + Sankey
----------------------------------

1. Extract noun phrases per column (situation, struggle, outcome).
2. Cluster them into themes.
3. Assign each row's situation/struggle/outcome to the closest theme.
4. Aggregate transitions into Sankey edges.

Outputs:
- <col>_keyphrases.csv
- <col>_themes_keyphrases.csv
- <col>_themes_assignments_keyphrases.csv
- sankey_edges.csv  (flows: situation → struggle → outcome)
"""

import threadpoolctl
threadpoolctl.threadpool_limits(limits=1, user_api="blas")
threadpoolctl.threadpool_limits(limits=1, user_api="openmp")

from sentence_transformers import SentenceTransformer
import argparse
import pandas as pd
import numpy as np
from collections import Counter
from pathlib import Path
import re, unicodedata, nltk, spacy
from nltk.stem import WordNetLemmatizer
from sklearn.cluster import KMeans
from sklearn.metrics.pairwise import cosine_distances

import torch
torch.set_num_threads(1)

import os
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"

try:
    import hdbscan
    HAS_HDBSCAN = True
except:
    HAS_HDBSCAN = False


# -------------------
# Preprocessing
# -------------------
def ensure_nltk():
    try:
        nltk.data.find("corpora/wordnet")
    except LookupError:
        nltk.download("wordnet", quiet=True)
    try:
        nltk.data.find("corpora/omw-1.4")
    except LookupError:
        nltk.download("omw-1.4", quiet=True)


def load_stopwords():
    return set("the and to a of in for on at with an is it that this so be as by or are was were from but have has had when i me my you your they their them we us our just really like get got make makes made much many very more less lot lots kind sort maybe".split())


def clean_text(text):
    if not isinstance(text, str):
        return ""
    text = text.lower()
    text = unicodedata.normalize("NFKD", text)
    text = re.sub(r"[^a-z\s]", " ", text)
    return re.sub(r"\s+", " ", text).strip()


lemmatizer = WordNetLemmatizer()
try:
    nlp = spacy.load("en_core_web_sm", disable=["ner"])
except OSError:
    import subprocess, sys
    subprocess.run([sys.executable, "-m", "spacy", "download", "en_core_web_sm"])
    nlp = spacy.load("en_core_web_sm", disable=["ner"])


def extract_noun_phrases(texts, stopwords, min_len=2, max_len=6):
    phrases = []
    for doc in nlp.pipe(texts, batch_size=64):
        for chunk in doc.noun_chunks:
            phrase = chunk.text.lower().strip()
            phrase = " ".join(
                [lemmatizer.lemmatize(tok) for tok in phrase.split() if tok not in stopwords]
            )
            if min_len <= len(phrase.split()) <= max_len:
                phrases.append(phrase)
    return phrases


# -------------------
# Embeddings + Clustering
# -------------------
# def embed_texts(texts, model):
#     return model.encode(texts, batch_size=64, show_progress_bar=False, normalize_embeddings=True)

def embed_texts(texts, model):
    if not texts:
        return np.zeros((0, model.get_sentence_embedding_dimension()))

    kwargs = dict(
        batch_size=64,
        show_progress_bar=False,
        normalize_embeddings=True,
        convert_to_numpy=True,
    )

    # add num_workers only if supported
    import inspect
    if "num_workers" in inspect.signature(model.encode).parameters:
        kwargs["num_workers"] = 0

    return model.encode(texts, **kwargs)


def cluster_kmeans(embeddings, k=8):
    n_samples = len(embeddings)

    if n_samples == 0:
        raise ValueError("No embeddings provided for clustering.")

    if n_samples < k:
        print(f"⚠️ Not enough samples ({n_samples}) for {k} clusters. "
              f"Reducing number of clusters to {n_samples}.")
        k = n_samples  # KMeans requires n_samples >= n_clusters

    km = KMeans(n_clusters=k, n_init="auto", random_state=42)
    labels = km.fit_predict(embeddings)
    return labels, km.cluster_centers_

def cluster_hdbscan(embeddings):
    if embeddings.shape[0] == 0:
        # no data at all
        return np.array([]), np.zeros((0, embeddings.shape[1])), {}

    clusterer = hdbscan.HDBSCAN(min_cluster_size=5, metric="euclidean")
    labels = clusterer.fit_predict(embeddings)
    centers = []
    cmap = {}

    unique_clusters = [c for c in set(labels) if c != -1]
    for c in sorted(unique_clusters):
        idx = (labels == c)
        centers.append(embeddings[idx].mean(axis=0))
        cmap[c] = len(centers) - 1

    if len(centers) == 0:
        # No clusters found — treat all as noise and use global mean
        centers = [embeddings.mean(axis=0)]
        cmap = {-1: 0}
        labels[:] = -1  # mark all as noise

    return labels, np.vstack(centers), cmap


def auto_label_themes(results, labels):
    df = pd.DataFrame([{"phrase": r[0], "count": r[1], "percent": r[2], "cluster": lab}
                       for r, lab in zip(results, labels)])
    clusters = sorted([c for c in df["cluster"].unique() if c != -1])
    themes = []
    label_map = {}
    for c in clusters:
        dfc = df[df["cluster"] == c]
        exemplar = dfc.sort_values("count", ascending=False)["phrase"].iloc[0]
        label_map[c] = exemplar
        themes.append({
            "theme_id": c,
            "theme_label": exemplar,
            "count": int(dfc["count"].sum()),
            "percent": round(dfc["count"].sum() / df["count"].sum() * 100, 2)
        })

    if not themes:
        # No clusters found (all noise)
        df["theme_label"] = "Noise"
        themes = [{
            "theme_id": -1,
            "theme_label": "Noise",
            "count": int(df["count"].sum()),
            "percent": 100.0
        }]
        label_map = {-1: "Noise"}
    else:
        df["theme_label"] = df["cluster"].map(label_map).fillna("Noise")

    return themes, df, label_map



# -------------------
# Main Analysis
# -------------------
def analyze_column(col_name, df_yes, model, stopwords, method="kmeans", num_themes=8, topn=50, min_count=2):
    texts = df_yes[col_name].dropna().tolist()
    cleaned = [clean_text(t) for t in texts]
    keyphrases = extract_noun_phrases(cleaned, stopwords)

    counter = Counter(keyphrases)
    items = [(ph, cnt) for ph, cnt in counter.items() if cnt >= min_count]
    items.sort(key=lambda x: x[1], reverse=True)
    items = items[:topn]
    total = sum(cnt for _, cnt in items) or 1
    results = [(ph, cnt, round(cnt / total * 100, 2)) for ph, cnt in items]

    if not results:
        return None, None, None

    phrases = [r[0] for r in results]
    emb = embed_texts(phrases, model)

    if method == "kmeans":
        labels, centers = cluster_kmeans(emb, num_themes)
        themes, assignments, label_map = auto_label_themes(results, labels)
    else:
        labels, centers, cmap = cluster_hdbscan(emb)
        themes, assignments, label_map = auto_label_themes(results, labels)

    return pd.DataFrame(themes), assignments, (centers, label_map)


def assign_rows_to_themes(df_yes, col_name, model, centers, label_map, stopwords):
    assigned = []
    for text in df_yes[col_name].fillna(""):
        phrases = extract_noun_phrases([clean_text(text)], stopwords)
        # if none extracted, assign "Noise"
        if not phrases:
            assigned.append("Noise")
            continue
        # look up clusters for each phrase (via embedding + nearest centroid)
        emb = embed_texts(phrases, model)
        dists = cosine_distances(emb, centers)
        nearest = dists.argmin(axis=1)
        # pick most frequent theme among phrases
        theme = Counter(nearest).most_common(1)[0][0]
        assigned.append(label_map.get(theme, "Noise"))
    return assigned




import plotly.graph_objects as go

def build_sankey(flows, out_file):
    # Get unique nodes per stage in order
    situations = sorted(flows["situation_theme"].unique())
    struggles = sorted(flows["struggle_theme"].unique())
    outcomes = sorted(flows["outcome_theme"].unique())

    nodes = situations + struggles + outcomes
    node_index = {n: i for i, n in enumerate(nodes)}

    sources, targets, values = [], [], []

    # situation → struggle
    for _, row in flows.iterrows():
        s = node_index[row["situation_theme"]]
        t = node_index[row["struggle_theme"]]
        sources.append(s)
        targets.append(t)
        values.append(row["count"])

    # struggle → outcome
    for _, row in flows.iterrows():
        s = node_index[row["struggle_theme"]]
        t = node_index[row["outcome_theme"]]
        sources.append(s)
        targets.append(t)
        values.append(row["count"])

    import plotly.graph_objects as go
    fig = go.Figure(data=[go.Sankey(
        node=dict(
            pad=15, thickness=20,
            line=dict(color="black", width=0.5),
            label=nodes,
            color=["#a6cee3"]*len(situations) + ["#b2df8a"]*len(struggles) + ["#fb9a99"]*len(outcomes)
        ),
        link=dict(source=sources, target=targets, value=values)
    )])

    fig.update_layout(title_text="JTBD Transformation Journey", font_size=12)
    fig.write_html(out_file)
    print(f"Saved Sankey diagram → {out_file}")

    


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--input", required=True)
    parser.add_argument("--cluster-method", choices=["kmeans", "hdbscan"], default="kmeans")
    parser.add_argument("--num-themes", type=int, default=8)
    args = parser.parse_args()

    ensure_nltk()
    stopwords = load_stopwords()
    model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

    df = pd.read_csv(args.input)

    input_path = Path(args.input)
    base_name = input_path.stem

    df_yes = df[df["jtbd"] == "YES"].drop_duplicates(subset=["comment"])

    results_by_col = {}
    centers_by_col = {}
    for col in ["situation", "struggle", "outcome"]:
        print(f"\n=== {col.upper()} ===")
        df_themes, assignments, (centers, label_map) = analyze_column(
            col, df_yes, model, stopwords, method=args.cluster_method, num_themes=args.num_themes
        )
        if df_themes is not None:
            results_by_col[col] = df_themes
            centers_by_col[col] = (centers, label_map)
            df_themes_sorted = df_themes.sort_values("count", ascending=False)
            print(df_themes_sorted[["theme_label", "count", "percent"]].head(10))

    # Assign each row to nearest theme in each column
    for col in ["situation", "struggle", "outcome"]:
        centers, label_map = centers_by_col[col]
        df_yes[f"{col}_theme"] = assign_rows_to_themes(
            df_yes, col, model, centers, label_map, stopwords
        )

    # -------------------------------------------------------
    # 🧹 Remove rows with unwanted discovered themes
    # -------------------------------------------------------
    remove_labels = {"grad school", "master s degree", "career path", "job market", "career goal",
                     "first job", "job prospect", "relevant experience", "higher pay",
                     "online resource"}

    before_count = len(df_yes)
    df_yes = df_yes[
        (~df_yes["situation_theme"].isin(remove_labels)) &
        (~df_yes["struggle_theme"].isin(remove_labels)) &
        (~df_yes["outcome_theme"].isin(remove_labels))
    ]
    after_count = len(df_yes)

    print(f"\nFiltered out {before_count - after_count} rows containing themes: {', '.join(remove_labels)}")


    # Aggregate flows
    flows = df_yes.groupby(["situation_theme", "struggle_theme", "outcome_theme"]).size().reset_index(name="count")

    # --- Sequential probabilities ---
    # P(struggle | situation)
    s2s = flows.groupby(["situation_theme", "struggle_theme"])["count"].sum().reset_index()
    s2s["P_struggle_given_situation"] = s2s.groupby("situation_theme")["count"].transform(lambda x: x / x.sum())

    # P(outcome | situation, struggle)
    sso = flows.copy()
    sso["P_outcome_given_situation_struggle"] = sso.groupby(["situation_theme", "struggle_theme"])["count"].transform(lambda x: x / x.sum())

    # Save probabilities
    s2s.to_csv(input_path.parent / f"{base_name}__transition_probs_situation_struggle.csv", index=False)
    sso.to_csv(input_path.parent / f"{base_name}__transition_probs_situation_struggle_outcome.csv", index=False)

    print("\nTop conditional transitions (P(outcome | situation, struggle)):\n")
    print(sso.sort_values("P_outcome_given_situation_struggle", ascending=False).head(15))

    # Print top flows in terminal
    flows_sorted = flows.sort_values("count", ascending=False)
    print("\nTop JTBD Journeys (situation → struggle → outcome):\n")
    for _, row in flows_sorted.head(20).iterrows():
        print(f"{row['situation_theme']} → {row['struggle_theme']} → {row['outcome_theme']}   ({row['count']})")

    # Export comments per flow if count above threshold
    threshold = 5  # you can make this an argument
    out_dir = Path(args.input).parent / "flow_comments"
    out_dir.mkdir(exist_ok=True)

    for _, row in flows.iterrows():
        s, st, o, count = row["situation_theme"], row["struggle_theme"], row["outcome_theme"], row["count"]

        if count >= threshold:
            subset = df_yes[
                (df_yes["situation_theme"] == s) &
                (df_yes["struggle_theme"] == st) &
                (df_yes["outcome_theme"] == o)
            ]

            # Save with a readable filename
            safe_name = f"{base_name}__{s}_{st}_{o}".replace(" ", "_").replace("/", "-")
            file_path = out_dir / f"{safe_name}.csv"
            subset[["situation", "struggle", "outcome", "comment"]].to_csv(file_path, index=False)

            print(f"Saved flow comments → {file_path}")


    # ============================================
    # 10️⃣ JTBD Archetype Discovery
    # ============================================
    print("\n=== JTBD ARCHETYPE DISCOVERY ===")

    # 1️⃣ Combine situation + struggle + outcome
    df_yes["journey_text"] = (
        df_yes["situation"].fillna("") + " | " +
        df_yes["struggle"].fillna("") + " | " +
        df_yes["outcome"].fillna("")
    ).str.strip()

    # 2️⃣ Embed full triplets
    journey_texts = df_yes["journey_text"].tolist()
    journey_emb = embed_texts(journey_texts, model)

    # 3️⃣ Cluster them
    if args.cluster_method == "kmeans":
        n_clusters = min(args.num_themes, len(df_yes))
        arche_labels, centers = cluster_kmeans(journey_emb, n_clusters)
    else:
        arche_labels, centers, cmap = cluster_hdbscan(journey_emb)

    df_yes["archetype_cluster"] = arche_labels

    # 4️⃣ Summarize archetypes
    from collections import Counter

    archetype_summaries = []
    for c in sorted(set(arche_labels)):
        if c == -1:
            continue
        subset = df_yes[df_yes["archetype_cluster"] == c]
        text_blob = " ".join(subset["journey_text"].tolist()).lower()
        words = [w for w in re.findall(r"[a-z]{3,}", text_blob) if w not in load_stopwords()]
        common = [w for w, _ in Counter(words).most_common(8)]
        archetype_summaries.append({
            "archetype_id": c,
            "n_examples": len(subset),
            "top_words": ", ".join(common),
            "example": subset["journey_text"].iloc[0]
        })

    df_archetypes = pd.DataFrame(archetype_summaries)
    df_archetypes.to_csv(input_path.parent / f"{base_name}__jtbd_archetypes.csv", index=False)

    print("\nTop Archetypes Discovered:\n")
    for _, row in df_archetypes.iterrows():
        arche_id = row["archetype_id"]
        subset = df_yes[df_yes["archetype_cluster"] == arche_id]

        # Get up to 3 examples per archetype
        examples = subset["journey_text"].dropna().head(3).tolist()
        # or: subset.sample(min(3, len(subset))) if you prefer random examples

        print(f"[{arche_id}] ({row['n_examples']} examples) → {row['top_words']}")
        for i, ex in enumerate(examples, 1):
            print(f"   ex{i}: {ex}...")
        print()

    # 5️⃣ Optional: Save per-archetype comments
    out_dir = Path(args.input).parent / "archetype_examples"
    out_dir.mkdir(exist_ok=True)
    for c in sorted(set(arche_labels)):
        if c == -1:
            continue
        df_yes[df_yes["archetype_cluster"] == c][
            ["situation", "struggle", "outcome", "comment"]
        ].to_csv(out_dir / f"{base_name}__archetype_{c}.csv", index=False)

    # ============================================
    # 3️⃣ Connect Archetypes to Behavioral Outcomes
    # ============================================
    print("\n=== ARCHETYPE → OUTCOME RELATIONSHIPS ===")

    # Ensure required columns exist
    if "archetype_cluster" in df_yes.columns and "outcome_theme" in df_yes.columns:
        # Compute joint counts and probabilities
        arche_outcome = (
            df_yes.groupby(["archetype_cluster", "outcome_theme"])
            .size()
            .reset_index(name="count")
        )

        # Normalize by archetype to get conditional probabilities
        arche_outcome["P_outcome_given_archetype"] = (
            arche_outcome.groupby("archetype_cluster")["count"]
            .transform(lambda x: x / x.sum())
        )

        # Sort for readability
        arche_outcome_sorted = (
            arche_outcome.sort_values(
                ["archetype_cluster", "P_outcome_given_archetype"],
                ascending=[True, False],
            )
        )

        # Save results
        out_path = input_path.parent / f"{base_name}__archetype_outcome_relationships.csv"
        arche_outcome_sorted.to_csv(out_path, index=False)
        print(f"Saved archetype-outcome relationships → {out_path}")

        # Pretty print top outcomes per archetype
        print("\nTop outcomes per archetype:\n")
        for c in sorted(df_yes["archetype_cluster"].unique()):
            subset = arche_outcome_sorted[arche_outcome_sorted["archetype_cluster"] == c]
            top = subset.head(3)
            if len(top) > 0:
                print(f"Archetype {c}:")
                for _, r in top.iterrows():
                    print(f"   {r['outcome_theme']} → {r['P_outcome_given_archetype']:.2f}")
                print()
    else:
        print("Missing required columns for archetype → outcome mapping.")

    # ============================================
    # CROSS-"SEGMENT" COMPARISON USING ARCHETYPES
    # ============================================
    print("\n=== CROSS-ARCHETYPE COMPARISON ===")

    # Treat each archetype as its own "segment"
    df_seg = df_yes[["archetype_cluster", "situation_theme", "struggle_theme", "outcome_theme"]].dropna()

    summary = []
    for archetype, group in df_seg.groupby("archetype_cluster"):
        top_sit = group["situation_theme"].value_counts(normalize=True).head(3)
        top_str = group["struggle_theme"].value_counts(normalize=True).head(3)
        top_out = group["outcome_theme"].value_counts(normalize=True).head(3)

        summary.append({
            "archetype": archetype,
            "top_situations": ", ".join([f"{s} ({p:.0%})" for s, p in top_sit.items()]),
            "top_struggles": ", ".join([f"{s} ({p:.0%})" for s, p in top_str.items()]),
            "top_outcomes": ", ".join([f"{s} ({p:.0%})" for s, p in top_out.items()])
        })

    df_summary = pd.DataFrame(summary)
    df_summary.to_csv(input_path.parent / f"{base_name}__cross_archetype_comparison.csv", index=False)
    print(df_summary.to_string(index=False))

    # ============================================
    # CROSS-SEGMENT COMPARISON BY OUTCOME
    # ============================================
    print("\n=== CROSS-OUTCOME COMPARISON ===")

    df_out = df_yes[["outcome_theme", "situation_theme", "struggle_theme"]].dropna()

    out_summary = []
    for outcome, group in df_out.groupby("outcome_theme"):
        top_sit = group["situation_theme"].value_counts(normalize=True).head(3)
        top_str = group["struggle_theme"].value_counts(normalize=True).head(3)
        out_summary.append({
            "outcome": outcome,
            "top_situations": ", ".join([f"{s} ({p:.0%})" for s, p in top_sit.items()]),
            "top_struggles": ", ".join([f"{s} ({p:.0%})" for s, p in top_str.items()])
        })

    df_out_summary = pd.DataFrame(out_summary)
    df_out_summary.to_csv(input_path.parent / f"{base_name}__cross_outcome_comparison.csv", index=False)
    print(df_out_summary.to_string(index=False))

    out_path = input_path.parent / f"{base_name}__jtbd_archetypes.csv"
    df_yes.to_csv(out_path, index=False)
    print(f"Saved archetype data → {out_path}")

     # ============================================
    # 🔹 JOB STATEMENT CLUSTERING (Top-level intents)
    # ============================================
    if "job_statement" in df_yes.columns:
        print("\n=== JOB STATEMENT CLUSTERING ===")

        job_texts = df_yes["job_statement"].dropna().unique().tolist()
        if len(job_texts) < 3:
            print("Not enough job statements to cluster.")
        else:
            job_emb = embed_texts(job_texts, model)

            # Pick a reasonable number of clusters (5–8 is good for intent themes)
            n_clusters = min(args.num_themes, len(job_texts))
            labels, centers = cluster_kmeans(job_emb, n_clusters)

            df_jobs = pd.DataFrame({
                "job_statement": job_texts,
                "job_cluster": labels
            })

            # Derive short auto-labels from top words in each cluster
            summaries = []
            for c in sorted(set(labels)):
                subset = df_jobs[df_jobs["job_cluster"] == c]
                blob = " ".join(subset["job_statement"].tolist()).lower()
                words = [w for w in re.findall(r"[a-z]{3,}", blob) if w not in load_stopwords()]
                common = [w for w, _ in Counter(words).most_common(8)]
                summaries.append({
                    "cluster": c,
                    "count": len(subset),
                    "top_words": ", ".join(common),
                    "example": subset["job_statement"].iloc[0][:200]
                })

            df_summary = pd.DataFrame(summaries)
            df_summary.to_csv(input_path.parent / f"{base_name}__job_statement_clusters.csv", index=False)
            print("\nTop Job Statement Clusters:\n")
            print(df_summary.to_string(index=False))

            # Add cluster labels back into df_yes
            df_map = df_jobs.set_index("job_statement")["job_cluster"].to_dict()
            df_yes["job_cluster"] = df_yes["job_statement"].map(df_map)

            # Save enhanced file
            out_path2 = input_path.parent / f"{base_name}__jtbd_with_job_clusters.csv"
            df_yes.to_csv(out_path2, index=False)
            print(f"Saved JTBD data with job clusters → {out_path2}")

if __name__ == "__main__":
    main()