"""Per-cluster SVD feature extraction."""

import numpy as np
import polars as pl
from sklearn.decomposition import TruncatedSVD


def cluster_svd_extract(
    lf: pl.LazyFrame,
    labels: list[int],
    feature_cols: list[str],
) -> dict[int, dict]:
    """Per-cluster SVD-based distinguishing feature extraction.

    For each cluster label:
    1. Extracts rows belonging to that cluster.
    2. Fits ``TruncatedSVD(n_components=min(N, len(feature_cols)))``.
    3. The absolute loadings on the first SVD component serve as
       *feature_strength* — higher values mean more distinguishing.

    Returns a mapping ``{cluster_label: {...}}`` where each value is::

        {
            'features': [...],              # top-10 feature names by strength
            'explained_variance_ratio': [...],  # per-component ratios
            'feature_strength': [...],      # per-feature loading magnitude
        }

    Parameters
    ----------
    lf:
        LazyFrame with rows aligned to *labels*.
    labels:
        Cluster labels (same length as *lf* row count, ``-1`` = noise).
    feature_cols:
        Numeric column names to analyse.

    Returns
    -------
    dict[int, dict]
        Per-cluster SVD extraction results.
    """
    available = [c for c in feature_cols if c in lf.collect_schema().names()]
    if not available:
        return {}

    df = lf.select(available).collect(streaming=True)
    mat = df.to_numpy().astype(np.float64)
    mat = np.nan_to_num(mat, nan=0.0, posinf=0.0, neginf=0.0)

    if len(labels) != len(mat):
        labels = labels[:len(mat)]

    unique_labels = sorted(set(labels))
    result: dict[int, dict] = {}

    for label in unique_labels:
        mask = np.array(labels) == label
        n = int(mask.sum())
        if n < 2:
            continue

        cluster_mat = mat[mask]
        n_comp = min(n, len(available))
        svd = TruncatedSVD(n_components=n_comp, random_state=42)
        svd.fit(cluster_mat)

        # First component loadings → feature strength
        if n_comp > 0 and svd.components_.shape[0] > 0:
            feature_strength = np.abs(svd.components_[0]).astype(np.float64)
        else:
            feature_strength = np.zeros(len(available), dtype=np.float64)

        total_strength = float(feature_strength.sum())
        if total_strength > 1e-10:
            feature_strength = feature_strength / total_strength

        top_indices = np.argsort(-feature_strength)
        top_k = min(10, len(available))

        result[int(label)] = {
            'features': [available[i] for i in top_indices[:top_k]],
            'explained_variance_ratio': svd.explained_variance_ratio_.tolist(),
            'feature_strength': feature_strength.tolist(),
        }

    return result