fn_registry/python/functions/datascience/project_clusters_2d.py

"""Proyeccion PCA-2D + KMeans sobre el mismo subset, con puntos y labels alineados.

Estandariza una sola vez las columnas numericas (z-score), proyecta a 2D con PCA
y clusteriza con KMeans sobre EXACTAMENTE la misma matriz escalada, de modo que
la proyeccion 2D (`points`) y la etiqueta de cluster (`labels`) quedan alineadas
fila a fila. Es la pieza que `pca_explained` + `kmeans_segments` no cubren: esas
dos estandarizan por separado y descartan los labels, asi que sus salidas no se
pueden cruzar para pintar un scatter PCA coloreado por cluster. Determinista.
"""

import math

import numpy as np
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA
from sklearn.metrics import silhouette_score
from sklearn.preprocessing import StandardScaler


def project_clusters_2d(
    columns: dict,
    k_min: int = 2,
    k_max: int = 8,
    max_points: int = 2000,
) -> dict:
    """Proyecta a 2D (PCA) y clusteriza (KMeans) el mismo subset estandarizado.

    PCA a 2D y KMeans se ajustan sobre la MISMA matriz estandarizada, por lo que
    `points` (proyeccion 2D) y `labels` (cluster por fila) quedan alineados por
    indice. El k se elige automaticamente por silhouette en el rango
    [k_min, min(k_max, n_rows-1)], igual criterio que `kmeans_segments`.
    Determinista: StandardScaler + PCA(random_state=0) + KMeans(random_state=0,
    n_init=10).

    Args:
        columns: mapa {nombre_columna: [valores numericos]}. Listas alineadas por
            fila (misma longitud). Columnas no numericas o con menos de 2 valores
            distintos se descartan. None/NaN marcan filas a descartar listwise
            (una fila se elimina si cualquier feature falta).
        k_min: numero minimo de clusters a probar (default 2).
        k_max: numero maximo de clusters a probar (default 8). Se acota a
            min(k_max, n_rows_validas-1).
        max_points: tope de puntos devueltos en `points`/`labels`. Si las filas
            usadas superan este tope, se submuestrea points y labels CONJUNTAMENTE
            con paso determinista para mantenerlos alineados. El fit (best_k,
            silhouette, centroides, perfiles) usa SIEMPRE todas las filas.

    Returns:
        dict con points (proyeccion 2D, posiblemente submuestreada a max_points),
        labels (cluster de cada point, alineado con points), centers_2d
        (centroides en espacio PCA, len == best_k), best_k, silhouette,
        explained_2d (varianza de PC1 y PC2), cluster_sizes (sobre n_used total),
        cluster_profiles (ver abajo), feature_names, n_used (filas del fit antes
        de muestreo) y note ("" si ok). Cada entrada de cluster_profiles:
        {cluster, size, pct, centroid_original (medias en escala original),
        centroid_z (z del centroide), distinctive (top 3 features por |z|)}.
        Con <2 columnas numericas o <max(3, k_min*2) filas validas devuelve
        best_k=0 y note "datos insuficientes" sin lanzar excepcion.
    """
    feature_names: list[str] = []

    def insufficient(names: list[str], n_used: int) -> dict:
        return {
            "best_k": 0,
            "points": [],
            "labels": [],
            "centers_2d": [],
            "cluster_profiles": [],
            "feature_names": names,
            "n_used": int(n_used),
            "note": "datos insuficientes",
        }

    try:
        if not isinstance(columns, dict) or not columns:
            return insufficient([], 0)

        # 1. Coerce a numerico, descartando columnas no parseables o constantes.
        numeric_cols: dict[str, list] = {}
        for name, values in columns.items():
            if not isinstance(values, (list, tuple)):
                continue
            coerced: list[float] = []
            usable = True
            for v in values:
                if v is None:
                    coerced.append(math.nan)
                    continue
                try:
                    coerced.append(float(v))
                except (TypeError, ValueError):
                    usable = False
                    break
            if not usable:
                continue
            # Menos de 2 valores distintos no aporta varianza -> descartar.
            distinct = {x for x in coerced if not math.isnan(x)}
            if len(distinct) < 2:
                continue
            numeric_cols[name] = coerced

        feature_names = list(numeric_cols.keys())
        if len(feature_names) < 2:
            return insufficient(feature_names, 0)

        # 2. Matriz alineada por fila + listwise deletion (cualquier NaN -> fuera).
        matrix = np.array(
            [numeric_cols[n] for n in feature_names], dtype=float
        ).T
        valid_mask = ~np.isnan(matrix).any(axis=1)
        data = matrix[valid_mask]

        n_used = int(data.shape[0])
        min_rows = max(3, k_min * 2)
        if n_used < min_rows:
            return insufficient(feature_names, n_used)

        # 3. Estandarizar UNA sola vez (guardamos el scaler para desestandarizar).
        scaler = StandardScaler()
        X_scaled = scaler.fit_transform(data)

        # 4. PCA a 2D sobre la matriz escalada.
        pca = PCA(n_components=2, random_state=0)
        pca.fit(X_scaled)
        proj = pca.transform(X_scaled)

        # 5. KMeans con seleccion automatica de k por silhouette (mismo X_scaled).
        upper_k = min(k_max, n_used - 1)
        if upper_k < k_min:
            return insufficient(feature_names, n_used)

        best = None  # (silhouette, k, model, labels)
        for k in range(k_min, upper_k + 1):
            model = KMeans(n_clusters=k, n_init=10, random_state=0)
            labels_k = model.fit_predict(X_scaled)
            if len(set(labels_k)) < 2:
                sil = -1.0
            else:
                sil = float(silhouette_score(X_scaled, labels_k))
            if best is None or sil > best[0]:
                best = (sil, k, model, labels_k)

        best_sil, best_k, best_model, labels = best

        # 6. Centroides KMeans (espacio escalado) proyectados al espacio PCA.
        centers_2d = pca.transform(best_model.cluster_centers_)

        # 7. Perfiles por cluster sobre TODAS las filas usadas.
        centroids_original = scaler.inverse_transform(best_model.cluster_centers_)
        cluster_sizes: list[int] = []
        cluster_profiles: list[dict] = []
        for c in range(best_k):
            size = int(np.sum(labels == c))
            cluster_sizes.append(size)
            z_vec = best_model.cluster_centers_[c]
            orig_vec = centroids_original[c]
            centroid_z = {
                feature_names[j]: float(z_vec[j]) for j in range(len(feature_names))
            }
            centroid_original = {
                feature_names[j]: float(orig_vec[j])
                for j in range(len(feature_names))
            }
            order = np.argsort(np.abs(z_vec))[::-1]
            distinctive = [feature_names[int(j)] for j in order[:3]]
            cluster_profiles.append(
                {
                    "cluster": int(c),
                    "size": size,
                    "pct": float(size / n_used) if n_used else 0.0,
                    "centroid_original": centroid_original,
                    "distinctive": distinctive,
                    "centroid_z": centroid_z,
                }
            )

        # 8. Muestreo determinista CONJUNTO de points + labels (mantiene alineacion).
        note = ""
        if n_used > max_points and max_points > 0:
            step = math.ceil(n_used / max_points)
            proj_out = proj[::step]
            labels_out = labels[::step]
            note = f"submuestreado a {len(proj_out)} de {n_used} puntos para visualizacion"
        else:
            proj_out = proj
            labels_out = labels

        points = [[float(row[0]), float(row[1])] for row in proj_out]
        labels_list = [int(v) for v in labels_out]
        centers_list = [[float(row[0]), float(row[1])] for row in centers_2d]
        explained_2d = [float(x) for x in pca.explained_variance_ratio_]

        return {
            "points": points,
            "labels": labels_list,
            "centers_2d": centers_list,
            "best_k": int(best_k),
            "silhouette": float(best_sil),
            "explained_2d": explained_2d,
            "cluster_sizes": cluster_sizes,
            "cluster_profiles": cluster_profiles,
            "feature_names": feature_names,
            "n_used": n_used,
            "note": note,
        }
    except Exception:
        # Lectura defensiva: nunca propagar excepciones al caller del EDA.
        return insufficient(feature_names, 0)