tianxuan/analysis/tools/clustering.py

"""Clustering handlers — HDBSCAN, KMeans, AgglomerativeClustering.

Contains:
- _handle_run_clustering: cluster selected columns
- _handle_filter_and_cluster: filter then auto-cluster
- _cluster_core: shared clustering pipeline extracted from both handlers
"""
import logging
import traceback
from analysis.constants import LOW_MEMORY_THRESHOLD_GB, LOW_MEMORY_THRESHOLD_ROWS, MAX_CLUSTER_FEATURES, MAX_SAMPLE_ROWS, RANDOM_SEED
from typing import Optional

import numpy as np
import polars as pl

from ..session_store import SessionStore, _generate_id
from ._helpers import (
    NUMERIC_DTYPES,
    _build_filter_expr,
    _resolve_dataset,
    _count_rows,
)

logger = logging.getLogger(__name__)
_clust_logger = logging.getLogger(__name__)


def _cluster_core(
    data: np.ndarray,
    feature_cols: list[str],
    algorithm: str,
    params: dict,
    random_state: int = RANDOM_SEED,
):
    """Shared clustering pipeline: scale, filter, reduce, fit, score.

    Args:
        data: Numpy array (n_samples, n_features).
        feature_cols: Column names matching data columns.
        algorithm: 'hdbscan', 'kmeans', or 'agglomerative'.
        params: Algorithm-specific parameters dict.
        random_state: Random seed.

    Returns:
        (labels_array, n_clusters_found, noise_count, quality_metrics, _ds_idx)
    """
    p = dict(params)

    # ── Low memory check ───────────────────────────────────────────────
    low_memory = False
    try:
        import psutil
        mem = psutil.virtual_memory()
        low_memory = mem.available < LOW_MEMORY_THRESHOLD_GB * 1024 ** 3  # 2 GB
    except ImportError:
        pass

    if low_memory and len(data) > LOW_MEMORY_THRESHOLD_ROWS:
        rng = np.random.default_rng(RANDOM_SEED)
        idx = rng.choice(len(data), MAX_SAMPLE_ROWS, replace=False)
        data = data[idx]

    # ── Downsampling to 50K max ────────────────────────────────────────
    _ds_idx = None
    if len(data) > MAX_SAMPLE_ROWS:
        rng = np.random.default_rng(RANDOM_SEED)
        _ds_idx = rng.choice(len(data), MAX_SAMPLE_ROWS, replace=False)
        data = data[_ds_idx]

    # ── Variance filtering ─────────────────────────────────────────────
    if data.shape[1] > 1:
        variances = np.var(data, axis=0)
        keep_var = np.where(variances >= 1e-10)[0]
        if len(keep_var) == 0:
            keep_var = np.arange(data.shape[1])
        if len(keep_var) < data.shape[1]:
            data = data[:, keep_var]
            feature_cols = [feature_cols[i] for i in keep_var]

    # Log clustering input
    _clust_logger.info(
        '[CLUSTER_INPUT] rows=%d cols=%d feature_cols=%s',
        len(data), data.shape[1] if len(data.shape) > 1 else 1, feature_cols,
    )

    # Standard scale
    from sklearn.preprocessing import StandardScaler
    _clust_logger.info('[CLUSTER_SCALE] Starting StandardScaler...')
    data_scaled = StandardScaler().fit_transform(data)
    _clust_logger.info('[CLUSTER_SCALE] Done. shape=%s', data_scaled.shape)

    # ── Correlation filtering ──────────────────────────────────────────
    if data_scaled.shape[1] > 1:
        corr = np.corrcoef(data_scaled.T)
        upper = np.triu(np.abs(corr), 1)
        to_drop = [j for j in range(upper.shape[1]) if any(upper[:, j] > 0.95)]
        if to_drop:
            keep = [j for j in range(data_scaled.shape[1]) if j not in to_drop]
            data_scaled = data_scaled[:, keep]
            feature_cols = [feature_cols[j] for j in keep]

    # ── SVD dimensionality reduction (if >50 features) ─────────────────
    n_features = data_scaled.shape[1]
    if n_features > MAX_CLUSTER_FEATURES:
        from sklearn.decomposition import TruncatedSVD
        svd = TruncatedSVD(n_components=MAX_CLUSTER_FEATURES, random_state=random_state)
        data_scaled = svd.fit_transform(data_scaled)

    # ── Hyperparameter auto-tuning ─────────────────────────────────────
    n = len(data_scaled)
    if algorithm == 'hdbscan':
        min_cluster = p.get('min_cluster_size', max(3, min(MAX_CLUSTER_FEATURES, n // 20)))
        p['min_cluster_size'] = min_cluster
        p['min_samples'] = p.get('min_samples', min_cluster)
        p['cluster_selection_epsilon'] = p.get('cluster_selection_epsilon', 0.0)
    else:
        p['min_cluster_size'] = min(
            p.get('min_cluster_size', 5),
            max(n // 2, 2),
        )

    # ── Model fit ──────────────────────────────────────────────────────
    from sklearn.cluster import HDBSCAN, MiniBatchKMeans, AgglomerativeClustering
    if algorithm == 'kmeans':
        n_clusters_param = p.get('n_clusters', 3)
        model = MiniBatchKMeans(
            n_clusters=n_clusters_param, random_state=random_state,
            batch_size=1024, n_init='auto',
        )
    elif algorithm in ('agglomerative', 'ward'):
        n_clusters_param = p.get('n_clusters', 5)
        model = AgglomerativeClustering(
            n_clusters=n_clusters_param,
            metric=p.get('metric', 'euclidean'),
            linkage=p.get('linkage', 'ward'),
            distance_threshold=None,
        )
    else:
        model = HDBSCAN(
            min_cluster_size=p['min_cluster_size'],
            min_samples=p.get('min_samples', None),
            metric=p.get('metric', 'euclidean'),
            cluster_selection_epsilon=p.get('cluster_selection_epsilon', 0.0),
        )

    labels_array = model.fit_predict(data_scaled)
    n_clusters_found = int(len(set(labels_array)) - (1 if -1 in labels_array else 0))
    noise_count = int((labels_array == -1).sum()) if algorithm == 'hdbscan' else 0
    noise_ratio = round(noise_count / len(labels_array), 4) if algorithm == 'hdbscan' else 0.0

    # ── Quality metrics ────────────────────────────────────────────────
    quality = {
        'n_clusters': n_clusters_found,
        'n_noise': noise_count,
        'noise_ratio': noise_ratio,
        'algorithm': algorithm,
    }

    if n_clusters_found > 1:
        try:
            from sklearn.metrics import silhouette_score
            score = silhouette_score(data_scaled, labels_array, random_state=random_state)
            quality['silhouette_score'] = round(float(score), 4)
        except Exception:
            logger.error("unknown failed: {}".format(traceback.format_exc()))
            quality['silhouette_score'] = None
        try:
            from sklearn.metrics import davies_bouldin_score
            dbs = davies_bouldin_score(data_scaled, labels_array)
            quality['davies_bouldin_score'] = round(float(dbs), 4)
        except Exception:
            logger.error("unknown failed: {}".format(traceback.format_exc()))
            quality['davies_bouldin_score'] = None
        try:
            from sklearn.metrics import calinski_harabasz_score
            chs = calinski_harabasz_score(data_scaled, labels_array)
            quality['calinski_harabasz_score'] = round(float(chs), 4)
        except Exception:
            logger.error("unknown failed: {}".format(traceback.format_exc()))
            quality['calinski_harabasz_score'] = None

    # Cluster sizes
    cluster_sizes = {}
    for label in np.unique(labels_array):
        cluster_sizes[int(label)] = int((labels_array == label).sum())
    quality['cluster_sizes'] = cluster_sizes

    return labels_array, n_clusters_found, noise_count, quality, _ds_idx


# ═══════════════════════════════════════════════════════════════════════
#  _handle_run_clustering
# ═══════════════════════════════════════════════════════════════════════

async def _handle_run_clustering(
    dataset_id: str,
    cluster_columns: list[str],
    algorithm: str = 'agglomerative',
    params: Optional[dict] = None,
    random_state: int = RANDOM_SEED,
) -> dict:
    """Run HDBSCAN, KMeans, or Agglomerative clustering on selected columns."""
    entry, err = _resolve_dataset(dataset_id)
    if err:
        return err

    lf: pl.LazyFrame = entry['lazyframe']
    p = params or {}

    # Collect cluster features - numeric only
    available = lf.collect_schema()
    available_names = available.names()
    available_dtypes = available.dtypes()
    feature_cols = [
        c for c in cluster_columns
        if c in available_names
        and available_dtypes[available_names.index(c)] in NUMERIC_DTYPES
    ]
    if not feature_cols:
        # Fall back: auto-select all numeric columns (exclude internal ones)
        feature_cols = [
            available_names[i] for i, dt in enumerate(available_dtypes)
            if dt in NUMERIC_DTYPES and not available_names[i].startswith('_')
        ][:10]
    if not feature_cols:
        # Fallback: try coercing Utf8 columns to Float64
        utf8_cols = [
            available_names[i] for i, dt in enumerate(available_dtypes)
            if dt == pl.Utf8 and not available_names[i].startswith('_')
        ][:10]
        if utf8_cols:
            for col in utf8_cols:
                try:
                    casted = lf.select(
                        pl.col(col).cast(pl.Float64, strict=False).alias(col)
                    ).collect(streaming=True)
                    null_ratio = casted[col].is_null().sum() / max(len(casted), 1)
                    if null_ratio < 0.9:
                        lf = lf.with_columns(
                            pl.col(col).cast(pl.Float64, strict=False)
                        )
                        feature_cols.append(col)
                except Exception:
                    logger.error("unknown failed: {}".format(traceback.format_exc()))
                    pass
        if not feature_cols:
            all_numeric = [f"{available_names[i]}" for i, dt in enumerate(available_dtypes)
                           if dt in NUMERIC_DTYPES]
            return {
                'error': f'No numeric cluster columns found in dataset. Available: {all_numeric}',
                'truncated': False
            }

    # Downsample before collect
    row_count = _count_rows(lf)
    if row_count > LOW_MEMORY_THRESHOLD_ROWS:
        lf = lf.collect(streaming=True).sample(n=LOW_MEMORY_THRESHOLD_ROWS, seed=RANDOM_SEED).lazy()

    df_features = lf.select(feature_cols).collect(streaming=True)

    # Drop all-NaN columns
    for col in df_features.columns:
        if df_features[col].is_null().all():
            df_features = df_features.drop(col)
            feature_cols = [c for c in feature_cols if c != col]

    if df_features.width == 0:
        return {'error': 'No valid feature columns after dropping all-NaN columns',
                'truncated': False}

    data = df_features.to_numpy()

    # Fill remaining NaN with column mean
    if data.dtype.kind == 'f':
        col_mean = np.nanmean(data, axis=0)
        col_mean = np.nan_to_num(col_mean, nan=0.0)
        data = np.where(np.isnan(data), col_mean, data)

    # Zero-sample guard
    if len(data) < 3:
        return {
            'cluster_result_id': _generate_id('clust'),
            'n_clusters': 0,
            'n_noise': 0,
            'quality_metrics': {'note': 'too few samples'},
            'truncated': False,
        }

    labels_array, n_found, noise_count, quality, _ds_idx = _cluster_core(
        data, feature_cols, algorithm, p, random_state,
    )


    result_id = _generate_id('clust')
    store = SessionStore()
    store.store_cluster_result(
        result_id=result_id,
        labels=labels_array.tolist(),
        n_clusters=n_found,
        params={
            'algorithm': algorithm,
            'params': p,
            'random_state': random_state,
            'feature_columns': feature_cols,
            '_ds_idx': _ds_idx.tolist() if _ds_idx is not None else None,
        },
        parent_dataset_id=dataset_id,
    )

    return {
        'cluster_result_id': result_id,
        'n_clusters': n_found,
        'n_noise': noise_count,
        'labels_sample': labels_array[:50].tolist(),
        'quality_metrics': quality,
    }


# ═══════════════════════════════════════════════════════════════════════
#  _handle_filter_and_cluster
# ═══════════════════════════════════════════════════════════════════════

async def _handle_filter_and_cluster(
    dataset_id: str,
    filters: list[dict],
    logic: str = 'and',
    algorithm: str = 'hdbscan',
    params: dict | None = None,
) -> dict:
    """Apply column-level filters then run clustering on the filtered result."""
    entry, err = _resolve_dataset(dataset_id)
    if err:
        return err

    lf: pl.LazyFrame = entry['lazyframe']
    full_schema = entry.get('schema', {})

    # ── Step 1: Apply filters ─────────────────────────────────────────────
    exprs = [_build_filter_expr(
        f['column'], f.get('op', 'eq'), f.get('value'),
    ) for f in filters]

    if exprs:
        if logic == 'and':
            combined = exprs[0]
            for e in exprs[1:]:
                combined = combined & e
        else:
            combined = exprs[0]
            for e in exprs[1:]:
                combined = combined | e
        filtered_lf = lf.filter(combined)
    else:
        filtered_lf = lf

    try:
        row_count = _count_rows(filtered_lf)
    except Exception:
        logger.error("unknown failed: {}".format(traceback.format_exc()))
        row_count = None


    # ── Step 2: Auto-detect numeric feature columns ───────────────────────
    available = filtered_lf.collect_schema()
    available_names = available.names()
    available_dtypes = available.dtypes()
    feature_cols = [
        available_names[i] for i, dt in enumerate(available_dtypes)
        if dt in NUMERIC_DTYPES and not available_names[i].startswith('_')
    ][:10]

    if not feature_cols:
        return {
            'error': 'No numeric columns found for clustering after filtering',
            'filtered_dataset_id': None,
            'truncated': False,
        }

    # ── Step 3: Store filtered dataset ────────────────────────────────────
    filtered_dataset_id = _generate_id('flt')
    store = SessionStore()
    store.store_dataset(
        dataset_id=filtered_dataset_id,
        lazyframe=filtered_lf,
        schema=full_schema,
        metadata={
            'row_count': row_count,
            'parent_dataset_id': dataset_id,
            'filter_count': len(filters),
            'logic': logic,
        },
        parent_id=dataset_id,
    )

    # ── Step 4: Run clustering ────────────────────────────────────────────
    p = params or {}
    total_count = _count_rows(filtered_lf)
    clamp_lf = filtered_lf
    if total_count > LOW_MEMORY_THRESHOLD_ROWS:
        clamp_lf = filtered_lf.collect(streaming=True).sample(n=LOW_MEMORY_THRESHOLD_ROWS, seed=RANDOM_SEED).lazy()

    df_features = clamp_lf.select(feature_cols).collect(streaming=True)

    for col in df_features.columns:
        if df_features[col].is_null().all():
            df_features = df_features.drop(col)
            feature_cols = [c for c in feature_cols if c != col]

    if df_features.width == 0:
        return {'error': 'No valid feature columns after dropping all-NaN columns',
                'filtered_dataset_id': filtered_dataset_id, 'truncated': False}

    data = df_features.to_numpy()

    if data.dtype.kind == 'f':
        col_mean = np.nanmean(data, axis=0)
        col_mean = np.nan_to_num(col_mean, nan=0.0)
        data = np.where(np.isnan(data), col_mean, data)

    if len(data) < 3:
        return {
            'cluster_result_id': _generate_id('clust'),
            'n_clusters': 0,
            'n_noise': 0,
            'quality_metrics': {'note': 'too few samples'},
            'filtered_dataset_id': filtered_dataset_id,
            'truncated': False,
        }

    labels_array, n_found, noise_count, quality, _ds_idx = _cluster_core(
        data, feature_cols, algorithm, p, RANDOM_SEED,
    )


    result_id = _generate_id('clust')
    store.store_cluster_result(
        result_id=result_id,
        labels=labels_array.tolist(),
        n_clusters=n_found,
        params={
            'algorithm': algorithm,
            'params': p,
            'random_state': RANDOM_SEED,
            'feature_columns': feature_cols,
            '_ds_idx': _ds_idx.tolist() if _ds_idx is not None else None,
        },
        parent_dataset_id=filtered_dataset_id,
    )

    return {
        'cluster_result_id': result_id,
        'filtered_dataset_id': filtered_dataset_id,
        'n_clusters': n_found,
        'n_noise': noise_count,
        'labels_sample': labels_array[:50].tolist(),
        'quality_metrics': quality,
        'feature_columns': feature_cols,
    }