import os
import sys
import pandas as pd
import numpy as np
from llm_client import call_gemini_embeddings
from sklearn.cluster import AgglomerativeClustering
from sklearn.metrics.pairwise import cosine_distances

from sklearn.metrics import pairwise_distances_argmin_min

def pick_representative(canonicals, cluster_embeddings):
    center_idx, _ = pairwise_distances_argmin_min(
        cluster_embeddings.mean(axis=0).reshape(1, -1),
        cluster_embeddings
    )
    return canonicals[center_idx[0]]

# --- Config ---
input_file = sys.argv[1]

# --- Load data ---
df = pd.read_csv(input_file)

required_cols = {"canonical", "post_url", "author", "comment", "quotes"}
missing = required_cols - set(df.columns)
if missing:
    raise ValueError(f"CSV must contain columns: {missing}")

texts = df['canonical'].dropna().unique().tolist()

# --- Step 1: Get embeddings (with caching) ---
out_emb_file = f'{input_file.split(".csv")[0]}_embeddings.parquet'

if os.path.exists(out_emb_file):
    print(f"Loading cached embeddings from {out_emb_file} ...")
    df_emb = pd.read_parquet(out_emb_file)
    texts = df_emb["canonical"].tolist()
    embeddings = np.array(df_emb["embedding"].tolist())
else:
    print(f"Embedding {len(texts)} canonicals...")
    batch_embeddings = call_gemini_embeddings(texts)
    if batch_embeddings is None or any(e is None for e in batch_embeddings):
        raise RuntimeError("Gemini embeddings failed for one or more canonicals")

    embeddings = np.array(batch_embeddings)

    # Save to parquet
    df_emb = pd.DataFrame({
        "canonical": texts,
        "embedding": embeddings.tolist()
    })
    df_emb.to_parquet(out_emb_file, index=False)
    print(f"Saved embeddings → {out_emb_file}")

# --- Step 2: Clustering ---
dist_matrix = cosine_distances(embeddings)

clusterer = AgglomerativeClustering(
    n_clusters=None,
    metric="precomputed",
    linkage="average",
    distance_threshold=0.25   # adjust: higher = looser groups
)
labels = clusterer.fit_predict(dist_matrix)

# --- Step 3: Output (canonical-level) ---
df_clusters = pd.DataFrame({
    "canonical": texts,
    "cluster": labels
})

# Cluster sizes (unique canonicals per cluster)
sizes = df_clusters["cluster"].value_counts().to_dict()
df_clusters["cluster_size"] = df_clusters["cluster"].map(sizes)

# Representative phrase per cluster (central canonical)
representatives = {}
for cid in np.unique(labels):
    idxs = [i for i, lbl in enumerate(labels) if lbl == cid]
    cluster_canonicals = [texts[i] for i in idxs]
    cluster_embeds = embeddings[idxs]
    representatives[cid] = pick_representative(cluster_canonicals, cluster_embeds)

df_clusters["representative"] = df_clusters["cluster"].map(representatives)

# --- Step 4: Add post variability ---
# Join back original data (with post_url, comment, quote)
df_with_posts = df_clusters.merge(
    df[["canonical", "post_url", "author", "comment", "quotes"]],
    on="canonical", how="left"
)

# Count unique posts per cluster
post_counts = df_with_posts.groupby("cluster")["post_url"].nunique().to_dict()
df_with_posts["unique_posts"] = df_with_posts["cluster"].map(post_counts)

# Sort clusters by size (and optionally by variability)
df_with_posts = df_with_posts.sort_values(
    ["cluster_size", "unique_posts", "cluster"], ascending=[False, False, True]
)

# --- Save ---
out_file = f'{input_file.split(".csv")[0]}_painpoint_clusters.csv'
df_with_posts.to_csv(out_file, index=False)
print(f"\nSaved grouped pain points → {out_file}")

# --- Print sample ---
print("\n=== Sample pain point groups (sorted by size) ===")
for cid, group in df_with_posts.groupby("cluster", sort=False):
    size = group["cluster_size"].iloc[0]
    posts = group["unique_posts"].iloc[0]
    repr_text = group["representative"].iloc[0]
    print(f"\nPain Point {cid} ({size} canonicals, {posts} unique posts) | Repr: {repr_text}")
    for ex in group["canonical"].head(5).tolist():
        print(" -", ex)