"""SVD noise removal (step 3 of preprocessing pipeline)."""

import numpy as np
import polars as pl


def svd_denoise(
    lf: pl.LazyFrame,
    feature_cols: list[str],
    variance_threshold: float = 0.95,
) -> tuple[pl.LazyFrame, dict]:
    """Remove noise from numeric features via TruncatedSVD.

    1. Fits TruncatedSVD on *feature_cols*.
    2. Keeps *k* components where cumulative explained variance
       ``trace(S[:k]) / trace(S) ≥ variance_threshold`` (default 0.95).
    3. Inverse-transforms to reconstruct the denoised matrix.
    4. **Replaces** the original *feature_cols* with denoised values.

    The noise components (removed by SVD) are returned as a dict for
    display, NOT as new feature columns.

    Parameters
    ----------
    lf:
        LazyFrame containing *feature_cols*.
    feature_cols:
        Numeric column names to denoise.  These columns are replaced
        in-place by their denoised reconstructions.
    variance_threshold:
        Fraction of total variance to preserve (0.0–1.0).
        Default 0.95 keeps 95 % of signal and treats 5 % as noise.

    Returns
    -------
    tuple[pl.LazyFrame, dict]
        ``(denoised_lf, noise_profile)`` where *noise_profile* contains:
        - ``kept_components`` : int — number of components kept
        - ``total_components`` : int — max possible components
        - ``kept_variance`` : float — fraction of variance kept
        - ``noise_components`` : list[dict] — per-component noise info
          (component index, explained variance, top-5 contributing features)
    """
    noise_profile = {
        'kept_components': 0,
        'total_components': 0,
        'kept_variance': 0.0,
        'noise_components': [],
    }

    available = [c for c in feature_cols if c in lf.collect_schema().names()]
    if not available:
        return lf, noise_profile

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

    n, d = mat.shape
    max_components = min(n, d - 1)
    if max_components < 1:
        return lf, noise_profile

    from sklearn.decomposition import TruncatedSVD
    svd = TruncatedSVD(n_components=max_components, random_state=42)
    _ = svd.fit_transform(mat)

    total_variance = float(np.sum(svd.explained_variance_))
    if total_variance <= 1e-10:
        return lf, noise_profile

    cumsum = np.cumsum(svd.explained_variance_)
    for k in range(1, max_components + 1):
        if cumsum[k - 1] / total_variance >= variance_threshold:
            break
    else:
        k = max_components

    # Re-fit with chosen k
    svd_k = TruncatedSVD(n_components=k, random_state=42)
    reduced = svd_k.fit_transform(mat)
    reconstructed = svd_k.inverse_transform(reduced)

    # --- Build noise profile (removed components) ---
    # The "noise" is the variance in components k..max_components
    kept_var_ratio = cumsum[k - 1] / total_variance if k > 0 else 0.0
    noise_profile = {
        'kept_components': k,
        'total_components': max_components,
        'kept_variance': round(float(kept_var_ratio), 4),
    }

    # Per-component noise info for removed components (k and beyond)
    noise_info = []
    for comp_idx in range(k, min(k + 10, max_components)):
        loading = np.abs(svd.components_[comp_idx])
        top5_idx = np.argsort(loading)[-5:][::-1]
        top5 = [
            {'feature': available[i], 'strength': round(float(loading[i]), 4)}
            for i in top5_idx
        ]
        noise_info.append({
            'component': int(comp_idx + 1),
            'explained_var_ratio': round(float(svd.explained_variance_ratio_[comp_idx]) * 100, 3),
            'top_features': top5,
        })
    noise_profile['noise_components'] = noise_info

    # --- Replace original columns with denoised values ---
    df_result = lf.collect(streaming=True)
    for i, col in enumerate(available):
        df_result = df_result.with_columns(
            pl.Series(col, reconstructed[:, i])
        )

    return df_result.lazy(), noise_profile