#!/usr/bin/env python3
"""Cluster `current_approach` and `hesitation` into themes, mirroring the
embedding + KMeans + LLM-label pattern already used for job statements in
jtbd_tag_cloud_noun_phrases.py.

These two fields carry the classic JTBD competitive forces that the pipeline
currently extracts but never analyzes:

  - current_approach  -> the "push" force: what the customer is currently
                         hiring to do the job (workarounds, hacks, incumbents).
  - hesitation        -> the "anxiety" force: what is stopping them switching.

For each field this script clusters the distinct non-empty texts with a local
SentenceTransformer model + silhouette-selected KMeans, labels each cluster
with the LLM (top-word fallback), and writes two files per field:

  <base>__current_approach_clusters.csv   cluster,label,count,top_words,example
  <base>__current_approach_themes.csv     comment,current_approach_theme
  <base>__hesitation_clusters.csv         cluster,label,count,top_words,example
  <base>__hesitation_themes.csv           comment,hesitation_theme

Usage:
    python extract_approach_and_hesitation_clusters.py <base>__jtbd_archetypes.csv
"""

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

import argparse
import re
import sys
from collections import Counter
from pathlib import Path

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

from sentence_transformers import SentenceTransformer
import numpy as np
import pandas as pd
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score

from llm_client import call_llm_chat_json

import torch
torch.set_num_threads(1)


STOPWORDS = 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()
)

FIELDS = ["current_approach", "hesitation"]


def embed_texts(texts, model):
    if not texts:
        return np.zeros((0, model.get_sentence_embedding_dimension()))
    return model.encode(texts, batch_size=64, show_progress_bar=False,
                        normalize_embeddings=True, convert_to_numpy=True)


def pick_k_via_silhouette(embeddings, max_k=8, floor_divisor=5, k_min=2):
    n_samples = len(embeddings)
    upper = min(max_k, n_samples - 1, max(k_min, n_samples // floor_divisor))
    if n_samples < 2 * k_min or upper < k_min:
        return max(1, min(n_samples, k_min))
    best_k, best_score = k_min, -1
    with threadpoolctl.threadpool_limits(limits=1):
        for k in range(k_min, upper + 1):
            km = KMeans(n_clusters=k, n_init="auto", random_state=42)
            labels = km.fit_predict(embeddings)
            if len(set(labels)) < 2:
                continue
            score = silhouette_score(embeddings, labels)
            if score > best_score:
                best_k, best_score = k, score
    return best_k


def cluster_kmeans(embeddings, k=8, floor_divisor=5):
    if len(embeddings) == 0:
        raise ValueError("No embeddings provided for clustering.")
    best_k = pick_k_via_silhouette(embeddings, max_k=k, floor_divisor=floor_divisor)
    if best_k != k:
        print(f"   -> auto-picked k={best_k} via silhouette (requested up to {k}, n={len(embeddings)})")
    with threadpoolctl.threadpool_limits(limits=1):
        km = KMeans(n_clusters=best_k, n_init="auto", random_state=42)
        labels = km.fit_predict(embeddings)
    return labels


def top_words_of(texts, n=8):
    blob = " ".join(texts).lower()
    words = [w for w in re.findall(r"[a-z]{3,}", blob) if w not in STOPWORDS]
    return ", ".join(w for w, _ in Counter(words).most_common(n))


def label_cluster_with_llm(kind, top_words, examples):
    system_message = (
        "You label clusters of customer research data for product analysis. Given "
        "the most frequent words and a few examples from a cluster, produce ONE "
        'short, specific label (2-5 words) naming the theme. Return ONLY this JSON: '
        '{"label": "..."}'
    )
    examples_block = "\n".join(f"{i + 1}. {ex}" for i, ex in enumerate(examples))
    user_message = f"Field: {kind}\nTop words: {top_words}\n\nExamples:\n{examples_block}"
    try:
        result = call_llm_chat_json(
            {"temperature": 0.2, "max_tokens": 60},
            system_message=system_message,
            user_message=user_message,
        )
    except Exception:
        result = None
    if result and isinstance(result.get("label"), str) and result["label"].strip():
        return result["label"].strip()
    return top_words.split(", ")[0].title() if top_words else "Unlabeled"


def cluster_field(df, field, model, out_dir, base_name):
    theme_col = f"{field}_theme"
    series = df[field].fillna("").astype(str)
    nonempty = series[series.str.strip() != ""]

    if nonempty.empty:
        print(f"  [{field}] no non-empty values; skipping.")
        return

    unique_texts = list(dict.fromkeys(nonempty.tolist()))  # dedupe, keep order
    print(f"  [{field}] {len(unique_texts)} distinct values (from {len(nonempty)} rows)")

    if len(unique_texts) < 3:
        print(f"  [{field}] too few distinct values to cluster (<3); assigning a single theme.")
        theme = top_words_of(unique_texts) or "Unlabeled"
        clusters = pd.DataFrame([{"cluster": 0, "label": theme, "count": len(unique_texts),
                                  "top_words": theme, "example": unique_texts[0][:200]}])
        text_to_theme = {t: theme for t in unique_texts}
    else:
        emb = embed_texts(unique_texts, model)
        labels = cluster_kmeans(emb, k=8, floor_divisor=5)

        members = {}
        for t, lab in zip(unique_texts, labels):
            members.setdefault(int(lab), []).append(t)

        text_to_theme = {}
        summaries = []
        for c in sorted(members):
            texts = members[c]
            tw = top_words_of(texts)
            label = label_cluster_with_llm(field, tw, texts[:6])
            summaries.append({"cluster": c, "label": label, "count": len(texts),
                              "top_words": tw, "example": texts[0][:200]})
            for t in texts:
                text_to_theme[t] = label

        clusters = (pd.DataFrame(summaries)
                    .sort_values("count", ascending=False)
                    .reset_index(drop=True))

    clusters.to_csv(out_dir / f"{base_name}__{field}_clusters.csv", index=False)
    print(f"  Saved clusters -> {base_name}__{field}_clusters.csv")
    print(clusters[["cluster", "label", "count"]].to_string(index=False))

    themes = df[["comment"]].copy()
    themes[theme_col] = series.map(lambda t: text_to_theme.get(t, "") if t.strip() else "")
    themes.to_csv(out_dir / f"{base_name}__{field}_themes.csv", index=False)
    print(f"  Saved per-row themes -> {base_name}__{field}_themes.csv")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("input", help="Path to <base>__jtbd_archetypes.csv")
    args = parser.parse_args()

    input_path = Path(args.input)
    if not input_path.exists():
        print(f"Error: input file not found: {input_path}")
        sys.exit(1)

    base_name = input_path.name.replace("__jtbd_archetypes.csv", "")
    out_dir = input_path.parent

    df = pd.read_csv(input_path)
    model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

    for field in FIELDS:
        if field not in df.columns:
            print(f"[{field}] column not present; skipping.")
            continue
        print(f"\n=== Clustering {field} ===")
        cluster_field(df, field, model, out_dir, base_name)

    print("\nDone.")


if __name__ == "__main__":
    main()