"""
Cluster confidence scoring for patient identification.

For each K-means cluster (in PCA-50d space), compute how well-separated
it is from its nearest neighbor.  Well-separated clusters are more likely
to represent real individual patients.
"""

import numpy as np
from sklearn.decomposition import PCA
from sklearn.cluster import KMeans
from scipy.spatial.distance import cdist


def compute_confidence(X, labels, patient_estimates, random_state=42):
    """Score each patient cluster by its separation from neighbors.

    Works in PCA-50d space (same space K-means operates in).

    Parameters
    ----------
    X : np.ndarray  (n_images, n_features)
        Full feature matrix (all classes together).
    labels : np.ndarray  (n_images,)
        Class label for each image.
    patient_estimates : dict
        class_name -> n_patients.
    random_state : int

    Returns
    -------
    cluster_scores : dict
        patient_id -> {
            "confidence": float,        # separation / internal_spread
            "intra_spread": float,      # mean distance to centroid
            "nearest_neighbor_dist": float,
            "nearest_neighbor_id": str,
            "images": list of filenames,  # set by caller
            "n_images": int,
        }
    """
    n_pca = min(50, X.shape[0] - 1, X.shape[1])
    pca = PCA(n_components=n_pca, random_state=random_state)
    X_pca = pca.fit_transform(X)

    cluster_scores = {}

    for class_name, n_clusters in patient_estimates.items():
        idx_class = np.where(labels == class_name)[0]
        X_class = X_pca[idx_class]

        kmeans = KMeans(n_clusters=n_clusters,
                        random_state=random_state, n_init=20)
        cluster_ids = kmeans.fit_predict(X_class)
        centroids = kmeans.cluster_centers_   # (n_clusters, n_pca)

        # Intra-cluster spread: mean distance from centroid
        intra_spreads = np.zeros(n_clusters)
        for k in range(n_clusters):
            mask = cluster_ids == k
            if mask.sum() > 0:
                dists = np.linalg.norm(X_class[mask] - centroids[k], axis=1)
                intra_spreads[k] = dists.mean()

        # Inter-cluster: distance to nearest OTHER centroid
        centroid_dists = cdist(centroids, centroids, metric='euclidean')
        np.fill_diagonal(centroid_dists, np.inf)
        nearest_dists = centroid_dists.min(axis=1)
        nearest_idx = centroid_dists.argmin(axis=1)

        # Confidence = separation / internal spread
        # (add epsilon to avoid division by zero)
        with np.errstate(divide='ignore', invalid='ignore'):
            confidences = nearest_dists / (intra_spreads + 1e-8)
        confidences = np.nan_to_num(confidences, nan=0.0)

        # Build result dict
        class_prefix = {
            "Bengin cases": "benign",
            "Malignant cases": "malig",
            "Normal cases": "normal",
        }[class_name]

        for k in range(n_clusters):
            pid = f"{class_prefix}_{k}"
            cluster_scores[pid] = {
                "confidence": float(confidences[k]),
                "intra_spread": float(intra_spreads[k]),
                "nearest_neighbor_dist": float(nearest_dists[k]),
                "nearest_neighbor_id": f"{class_prefix}_{int(nearest_idx[k])}",
                "class_name": class_name,
                "cluster_id": k,
            }

    return cluster_scores, X_pca, pca