from collections import defaultdict

from sentence_transformers import SentenceTransformer
from sklearn.cluster import KMeans
from sklearn.metrics import adjusted_rand_score, silhouette_score


texts = [
    "Reset the customer account password.",
    "Email the customer account password reset link.",
    "Review the customer account recovery settings.",
    "Build the website release image.",
    "Deploy the website release image.",
    "Roll back the website release image.",
    "Create the database backup archive.",
    "Restore the database backup archive.",
    "Verify the database backup archive.",
]
expected_theme_ids = [0, 0, 0, 1, 1, 1, 2, 2, 2]

model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device="cpu")
embeddings = model.encode(
    texts,
    normalize_embeddings=True,
    show_progress_bar=False,
)

clusterer = KMeans(n_clusters=3, random_state=7, n_init=10)
labels = clusterer.fit_predict(embeddings)

groups = defaultdict(list)
for text, label in zip(texts, labels):
    groups[int(label)].append(text)

for cluster_number, label in enumerate(
    sorted(groups, key=lambda key: groups[key][0]),
    start=1,
):
    print(f"Cluster {cluster_number}:")
    for text in groups[label]:
        print(f"- {text}")
    print()

score = silhouette_score(embeddings, labels)
agreement = adjusted_rand_score(expected_theme_ids, labels)
print(f"Silhouette score: {score:.2f}")
print(f"Adjusted Rand index: {agreement:.2f}")

if score < 0.40 or agreement < 1.0:
    raise SystemExit("sample clustering checks failed")