import os
import pandas as pd
import numpy as np
import hdbscan
from openai import OpenAI
import sys

# 🔑 Set your API key (or export OPENAI_API_KEY in your shell)
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# --- Load data ---
df = pd.read_csv(sys.argv[1])
canonicals = df['canonical'].dropna().unique().tolist()

# --- Get embeddings from OpenAI ---
print("Embedding", len(canonicals), "canonical pain points...")
embeddings = []
batch_size = 100  # keep requests small
for i in range(0, len(canonicals), batch_size):
    batch = canonicals[i:i+batch_size]
    resp = client.embeddings.create(
        model="text-embedding-3-small",
        input=batch
    )
    for e in resp.data:
        embeddings.append(e.embedding)

embeddings = np.array(embeddings)

# --- Cluster with HDBSCAN ---
clusterer = hdbscan.HDBSCAN(min_cluster_size=5, metric="euclidean")
labels = clusterer.fit_predict(embeddings)

# --- Build DataFrame with cluster assignments ---
clustered_df = pd.DataFrame({
    "canonical": canonicals,
    "cluster": labels
})

# Merge cluster labels back into original dataset
df_with_clusters = df.merge(clustered_df, on="canonical", how="left")

# --- Summaries ---
cluster_counts = df_with_clusters[df_with_clusters["cluster"] != -1]["cluster"].value_counts()
sample_per_cluster = clustered_df.groupby("cluster")["canonical"].apply(
    lambda x: x.sample(min(5, len(x)), random_state=1).tolist()
)

print("\n=== Total Comment Counts per Cluster ===")
print(cluster_counts)

print("\n=== Sample Canonicals per Cluster ===")
for cluster_id, examples in sample_per_cluster.items():
    print(f"\nCluster {cluster_id}:")
    for ex in examples:
        print(" -", ex)

# Save dataset with cluster labels

clustered_pain_points_file = '%s_with_clusters.csv' % sys.argv[1].split('.csv')[0]
df_with_clusters.to_csv(clustered_pain_points_file, index=False)
print("\nSaved results to %s" % clustered_pain_points_file)

from sklearn.cluster import KMeans
import numpy as np

# --- Step 1: Find largest cluster (ignore noise -1) ---
cluster_counts = df_with_clusters[df_with_clusters["cluster"] != -1]["cluster"].value_counts()
largest_cluster_id = cluster_counts.idxmax()
print(f"Largest cluster is {largest_cluster_id} with {cluster_counts.max()} comments")

# --- Step 2: Get canonicals in that cluster ---
largest_canonicals = clustered_df[clustered_df["cluster"] == largest_cluster_id]["canonical"].tolist()

# --- Step 3: Embed canonicals ---
embeddings = []
batch_size = 100
for i in range(0, len(largest_canonicals), batch_size):
    batch = largest_canonicals[i:i+batch_size]
    resp = client.embeddings.create(
        model="text-embedding-3-small",
        input=batch
    )
    for e in resp.data:
        embeddings.append(e.embedding)

embeddings = np.array(embeddings)

# --- Step 4: Subcluster with KMeans ---
kmeans = KMeans(n_clusters=4, random_state=42)  # try 3–6
sub_labels = kmeans.fit_predict(embeddings)

subclustered_df = pd.DataFrame({
    "canonical": largest_canonicals,
    "subcluster": sub_labels
})

# --- Step 5: Summarize ---
print("\n=== Canonical Subcluster Counts (Largest Cluster) ===")
print(subclustered_df["subcluster"].value_counts())

print("\n=== Sample Canonicals per Subcluster ===")
for cid in subclustered_df["subcluster"].value_counts().index:
    examples = subclustered_df[subclustered_df["subcluster"] == cid]["canonical"].head(5).tolist()
    print(f"\nSubcluster {cid}:")
    for ex in examples:
        print(" -", ex)

from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

# --- Step 1: Get embeddings for all comments (reuse if already computed) ---
text_df = df.copy()
text_df["text"] = text_df["quotes"].fillna(text_df["comment"])
text_df = text_df.dropna(subset=["text"])  # remove rows where both are null

texts = text_df["text"].tolist()
embeddings = []

batch_size = 100
for i in range(0, len(texts), batch_size):
    batch = texts[i:i+batch_size]
    resp = client.embeddings.create(
        model="text-embedding-3-small",
        input=batch
    )
    for e in resp.data:
        embeddings.append(e.embedding)

embeddings = np.array(embeddings)

# --- Step 2: Embed your "search query" ---
query = "Development Tools Frustration:* The build system and tools associated with Zephyr, notably Kconfig and device tree overlays, are frequently cited sources of frustration. Borrowed from Linux, these systems demand a nuanced understanding that can deter newcomers and seasoned developers alike. Despite their utility in configuring kernels to support specific features, the complexity of these tools often overshadows their benefits. Users advocate for enhanced tooling that simplifies these processes, suggesting that improvements could significantly lower the barrier of entry for adopting Zephyr in various development environments."
query_embedding = client.embeddings.create(
    model="text-embedding-3-large",
    input=[query]
).data[0].embedding

# --- Step 3: Compute similarity ---
sims = cosine_similarity([query_embedding], embeddings)[0]

# --- Step 4: Get top results ---
top_idx = sims.argsort()[-100:][::-1]  # top 100 matches
emotional_related = [texts[i] for i in top_idx]

print("\n=== Top 100 Matches for Query ===")
for i in top_idx:
    row = text_df.iloc[i]
    similarity_score = sims[i]  # how close this text is to the query
    print(f"\n- Text: {row['text']}")
    print(f"  Quote: {row['quotes']}")
    print(f"  Comment: {row['comment']}")
    print(f"  Similarity: {similarity_score:.4f}")


# # === TOPIC MODELING cu BERTopic (pe toate canonical-urile) ===
# from sklearn.feature_extraction.text import CountVectorizer
# try:
#     from bertopic import BERTopic
#     HAS_BERTOPIC = True
# except Exception as e:
#     print("BERTopic nu este instalat. Trec pe fallback NMF. Detalii:", e)
#     HAS_BERTOPIC = False

# if HAS_BERTOPIC:
#     # Avem deja: `canonicals` (list[str]) și `embeddings` (np.array shape [N, D]) din codul tău anterior.
#     # Dacă 'embeddings' a fost suprascris pentru subclusterul mare, recreăm embedding-urile pt. TOATE canonicals:
#     def embed_texts(texts, batch_size=100):
#         out = []
#         for i in range(0, len(texts), batch_size):
#             batch = texts[i:i+batch_size]
#             r = client.embeddings.create(model="text-embedding-3-small", input=batch)
#             out.extend([e.embedding for e in r.data])
#         return np.array(out)

#     # Asigură-te că 'embeddings' corespunde lui 'canonicals' (toate). Dacă ai dubii, re-embed:
#     if embeddings.shape[0] != len(canonicals):
#         print("Re-embed pentru TOATE canonicals, ca să corespundă BERTopic...")
#         embeddings_all = embed_texts(canonicals)
#     else:
#         embeddings_all = embeddings

#     # Vectorizer mai „strict” ca să curețe zgomotul
#     vectorizer = CountVectorizer(ngram_range=(1, 2), min_df=2, stop_words="english")

#     # Construim modelul
#     topic_model = BERTopic(
#         language="english",
#         vectorizer_model=vectorizer,
#         calculate_probabilities=True,
#         low_memory=True,           # util pe seturi mai mari
#         nr_topics=None             # lasă modelul să decidă; poți seta un număr dacă vrei
#     )

#     print("\n[BERTopic] Antrenez topic model...")
#     topics, probs = topic_model.fit_transform(canonicals, embeddings_all)

#     # Info topicuri (id, Count, Name)
#     topic_info = topic_model.get_topic_info()
#     print("\n=== Topic Info (primele 10) ===")
#     print(topic_info.head(10))

#     # Top cuvinte per topic
#     def pretty_topic_terms(topic_id, topn=10):
#         terms = topic_model.get_topic(topic_id)
#         if terms is None:
#             return []
#         return [t[0] for t in terms[:topn]]

#     # DataFrame final: canonical, topic_id, topic_name (label generat)
#     # topic_model._update_topic_labels() creează etichete scurte din top terms
#     topic_model.set_topic_labels(topic_model.generate_topic_labels(nr_words=3))
#     labels_map = dict(zip(topic_info["Topic"], topic_info["Name"]))

#     # Dacă vrei probabilitatea maximă per document:
#     if probs is not None:
#         max_probs = probs.max(axis=1)  # vector [num_docs]
#     else:
#         max_probs = [None] * len(canonicals)

#     df_topics_all = pd.DataFrame({
#         "canonical": canonicals,
#         "topic_id": topics,
#         "topic_name": [labels_map.get(t, f"Topic {t}") for t in topics],
#         "topic_prob_max": max_probs
#     })

#     # (Opțional) filtrare: doar cel mai mare cluster HDBSCAN
#     df_topics_largest = None
#     if "cluster" in df_with_clusters.columns:
#         df_tmp = pd.DataFrame({"canonical": canonicals, "topic_id": topics})
#         df_topics_merged = df_tmp.merge(df_with_clusters[["canonical", "cluster"]], on="canonical", how="left")
#         if (df_topics_merged["cluster"] != -1).any():
#             largest_cluster_id = df_with_clusters[df_with_clusters["cluster"] != -1]["cluster"].value_counts().idxmax()
#             mask = df_topics_merged["cluster"] == largest_cluster_id
#             ids_in_largest = set(df_topics_merged.loc[mask, "canonical"])
#             df_topics_largest = df_topics_all[df_topics_all["canonical"].isin(ids_in_largest)].copy()

#     # Print sumar scurt
#     print("\n=== Distribuție topicuri (toate canonical-urile) ===")
#     print(df_topics_all["topic_name"].value_counts().head(15))

#     if df_topics_largest is not None:
#         print("\n=== Distribuție topicuri în CEL MAI MARE CLUSTER (HDBSCAN) ===")
#         print(df_topics_largest["topic_name"].value_counts().head(15))

#     # Salvează rezultate
#     out_all = f'{sys.argv[1].split(".csv")[0]}_topics_all.csv'
#     df_topics_all.to_csv(out_all, index=False)
#     print(f"\n[BERTopic] Salvat: {out_all}")

#     if df_topics_largest is not None:
#         out_largest = f'{sys.argv[1].split(".csv")[0]}_topics_largest_cluster.csv'
#         df_topics_largest.to_csv(out_largest, index=False)
#         print(f"[BERTopic] Salvat: {out_largest}")

#     # (Opțional) extrage top termeni per topic și salvează
#     rows = []
#     for tid in topic_info["Topic"].tolist():
#         if tid == -1:  # -1 = outliers în BERTopic
#             continue
#         terms = pretty_topic_terms(tid, topn=10)
#         rows.append({"topic_id": tid, "topic_name": labels_map.get(tid, f"Topic {tid}"), "top_terms": ", ".join(terms)})
#     df_terms = pd.DataFrame(rows).sort_values("topic_id")
#     out_terms = f'{sys.argv[1].split(".csv")[0]}_topic_terms.csv'
#     df_terms.to_csv(out_terms, index=False)
#     print(f"[BERTopic] Salvat top terms: {out_terms}")
