metrics.py

"""
Metrics evaluation script for clustering analysis.
Converts embeddings to clusters and evaluates clustering quality.
"""

import asyncio
import hashlib
from pathlib import Path
from typing import List

import numpy as np
import pandas as pd
from infinity_client import Client
from infinity_client.api.default import embeddings
from infinity_client.models import OpenAIEmbeddingInputText
from metric_learn import ITML
from sklearn.cluster import AgglomerativeClustering
from sklearn.metrics import (
    calinski_harabasz_score,
    davies_bouldin_score,
    silhouette_score,
)
from sklearn.metrics.pairwise import cosine_distances

# Configuration
INPUT_FILE = "input.xlsx"
DISTANCE_THRESHOLD = 0.5  # порог кластеризации
EMB_DIM = 256

INFINITY_BASE_URL = "https://api.innohassle.ru/infinity"
MODEL_NAME = "sergeyzh/BERTA"
EMB_DIR = Path(".emb_cache")
EMB_DIR.mkdir(exist_ok=True)

client = Client(base_url=INFINITY_BASE_URL)


def get_cache_key(text: str, dim: int) -> str:
    raw = f"{MODEL_NAME}_{dim}_{text}"
    return hashlib.sha256(raw.encode()).hexdigest()


def emb_path(key: str) -> Path:
    return EMB_DIR / f"{key}.npy"


def load_emb(key: str) -> np.ndarray | None:
    p = emb_path(key)
    if p.exists():
        try:
            return np.load(p)
        except Exception:
            return None
    return None


def save_emb(key: str, arr: np.ndarray):
    np.save(emb_path(key), arr.astype(np.float32))


async def generate_embeddings(
    texts: List[str],
    batch_size: int = 32,
    normalize: bool = True,
    truncate_dim: int = 256,
    prefix: str = "categorize_entailment: ",
) -> np.ndarray:
    texts_pref = [prefix + t for t in texts]

    all_embs = [None] * len(texts_pref)
    to_fetch = []

    # cache
    for i, t in enumerate(texts_pref):
        key = get_cache_key(t, truncate_dim)
        loaded = load_emb(key)
        if loaded is not None:
            all_embs[i] = loaded
        else:
            to_fetch.append((i, t, key))

    # fetch
    if to_fetch:
        batches = [to_fetch[i : i + batch_size] for i in range(0, len(to_fetch), batch_size)]

        for batch in batches:
            texts_batch = [t for _, t, _ in batch]

            response = await embeddings.asyncio(
                client=client,
                body=OpenAIEmbeddingInputText(
                    model=MODEL_NAME,
                    input_=texts_batch,
                    dimensions=truncate_dim,
                ),
            )

            for (idx, _txt, key), item in zip(batch, response.data):
                arr = np.asarray(item.embedding, dtype=np.float32)
                all_embs[idx] = arr
                save_emb(key, arr)

    # normalize
    if normalize:
        for i, v in enumerate(all_embs):
            n = np.linalg.norm(v)
            if n > 0:
                all_embs[i] = v / n

    return np.vstack(all_embs)


def cluster_compactness(embeddings, labels):
    uniq = np.unique(labels)
    vals = []

    for lbl in uniq:
        idx = np.where(labels == lbl)[0]
        if len(idx) < 2:
            vals.append(0)
            continue
        cluster = embeddings[idx]
        centroid = np.mean(cluster, axis=0, keepdims=True)
        vals.append(cosine_distances(cluster, centroid).mean())
    return float(np.mean(vals))


def cluster_separation(embeddings, labels):
    uniq = np.unique(labels)
    cents = []
    for lbl in uniq:
        idx = np.where(labels == lbl)[0]
        cents.append(np.mean(embeddings[idx], axis=0))
    cents = np.vstack(cents)
    D = cosine_distances(cents)
    mask = ~np.eye(len(cents), dtype=bool)
    return float(D[mask].mean())


def cluster_stats(labels):
    labels = np.array(labels)
    uniq = np.unique(labels)
    sizes = [np.sum(labels == u) for u in uniq]
    return dict(
        n_clusters=len(uniq),
        min_cluster_size=int(np.min(sizes)),
        max_cluster_size=int(np.max(sizes)),
        mean_cluster_size=float(np.mean(sizes)),
        median_cluster_size=float(np.median(sizes)),
    )




def train_itml_metric(X_all, must_link_pairs, cannot_link_pairs):
    # Convert to arrays with shape (n, 2) even if empty
    ml = np.asarray(must_link_pairs, dtype=int).reshape(-1, 2)
    cl = np.asarray(cannot_link_pairs, dtype=int).reshape(-1, 2)

    n_ml, n_cl = len(ml), len(cl)

    # CASE 1: no constraints at all -> return None (identity metric)
    if n_ml == 0 and n_cl == 0:
        return None

    # Build pairs + labels for ITML
    if n_ml > 0 and n_cl > 0:
        pairs_idx = np.vstack([ml, cl])
        y_pairs = np.hstack(
            [
                np.ones(n_ml, dtype=int),  # must-link
                -np.ones(n_cl, dtype=int),  # cannot-link
            ]
        )
    elif n_ml > 0:  # only must-link
        pairs_idx = ml
        y_pairs = np.ones(n_ml, dtype=int)
    else:  # only cannot-link
        pairs_idx = cl
        y_pairs = -np.ones(n_cl, dtype=int)

    itml = ITML(preprocessor=X_all, verbose=False)
    itml.fit(pairs_idx, y_pairs)
    return itml


async def main():
    # Load data
    print("Loading data...")
    df = pd.read_excel(INPUT_FILE)
    df["название сте"] = df["название сте"].astype(str)
    print(f"Loaded {len(df)} rows")

    # Generate embeddings
    print("Generating embeddings...")
    embeddings_array = await generate_embeddings(
        df["название сте"].tolist(),
        truncate_dim=EMB_DIM,
    )
    print(f"Embeddings shape: {embeddings_array.shape}")

    # Clustering
    print("Performing clustering...")

    # Define constraints (empty by default - can be populated from external source)
    must_link_pairs = []  # List of (idx1, idx2) tuples
    cannot_link_pairs = []  # List of (idx1, idx2) tuples

    # Use ITML if constraints exist, otherwise use regular clustering
    itml = train_itml_metric(
        embeddings_array,
        must_link_pairs,
        cannot_link_pairs,
    )

    print("Training ITML metric with constraints...")
    print("Transforming embeddings with ITML...")
    embeddings_transformed = itml.transform(embeddings_array) if itml is not None else embeddings_array

    clusterer = AgglomerativeClustering(
        n_clusters=None,
        distance_threshold=DISTANCE_THRESHOLD,
        linkage="ward",
        metric="euclidean",
    )
    labels = clusterer.fit_predict(embeddings_transformed)

    df["cluster_id"] = labels
    print(f"Clustering complete. Number of clusters: {len(np.unique(labels))}")

    # Compute metrics
    print("\n=== METRICS ===")

    silhouette = silhouette_score(embeddings_array, labels, metric="cosine")
    print(f"Silhouette: {silhouette:.4f}")

    db_index = davies_bouldin_score(embeddings_array, labels)
    print(f"Davies-Bouldin: {db_index:.4f}")

    ch_score = calinski_harabasz_score(embeddings_array, labels)
    print(f"Calinski-Harabasz: {ch_score:.4f}")

    compactness = cluster_compactness(embeddings_array, labels)
    print(f"Compactness: {compactness:.4f}")

    separation = cluster_separation(embeddings_array, labels)
    print(f"Separation: {separation:.4f}")

    print("\nStats:")
    stats = cluster_stats(labels)
    for k, v in stats.items():
        print(f"  {k}: {v}")

    return df, labels, embeddings_array


if __name__ == "__main__":
    asyncio.run(main())