"""StandardScaler normalisation (step 2 of preprocessing pipeline)."""

import numpy as np
import polars as pl
from sklearn.preprocessing import StandardScaler


def normalize_features(lf: pl.LazyFrame, feature_cols: list[str]) -> pl.LazyFrame:
    """Apply StandardScaler normalisation to numeric feature columns.

    Collects *feature_cols* into a numpy array, replaces ``NaN`` with
    zero, centres to μ=0 σ=1 via ``sklearn.preprocessing.StandardScaler``,
    then returns the result as *lf* extended with ``_norm_*`` prefixed
    columns.

    Parameters
    ----------
    lf:
        LazyFrame containing *feature_cols*.
    feature_cols:
        Numeric column names to normalise.  Columns not present in
        *lf* are silently skipped; non-numeric columns raise an error.

    Returns
    -------
    pl.LazyFrame
        Input frame with additional ``_norm_{col}`` columns.

    Memory
    ------
    Peak memory ≈ collected numpy matrix + scaler workspace
    (well under the 800 MB budget for typical inputs).
    """
    available = [c for c in feature_cols if c in lf.collect_schema().names()]
    if not available:
        return lf

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

    scaler = StandardScaler()
    scaled = scaler.fit_transform(mat)

    # Build new columns with _norm_ prefix
    df_result = lf.collect(streaming=True)
    for i, col in enumerate(available):
        df_result = df_result.with_columns(
            pl.Series(f'_norm_{col}', scaled[:, i])
        )

    return df_result.lazy()
