"""
IP distance functions for Polars LazyFrames.

Provides subnet-mask distance and Haversine geographic distance
between pairs of IP columns.  All distances are normalised to [0, 1];
invalid / non-geolocatable IPs produce ``NaN``.

Typical usage::

    from analysis.distance import ip_distance

    lf = pl.scan_csv("flows.csv")
    lf = ip_distance(lf, ["src_ip", "dst_ip"])
    lf = ip_distance(lf, ["proxy_ip", "origin_ip"], subnet_masks=[16, 24])
"""

from __future__ import annotations

import math
from typing import Optional

import polars as pl

from analysis.geoip import ip_to_int, lookup

# ---------------------------------------------------------------------------
#  Earth constant
# ---------------------------------------------------------------------------

_HALF_EARTH_CIRCUMFERENCE_KM = 20_037.508  # equatorial / 2


# ---------------------------------------------------------------------------
#  Core distance helpers
# ---------------------------------------------------------------------------

def _subnet_mask_distance(ip1: Optional[str], ip2: Optional[str], mask: int) -> float:
    """Normalised subnet-mask distance between two IP strings.

    Computes ``abs(int(ip1)/2**mask - int(ip2)/2**mask) / 2**(32-mask)``.
    Returns ``NaN`` when either IP is ``None`` or unparseable.
    """
    if ip1 is None or ip2 is None:
        return float("nan")
    if not isinstance(ip1, str):
        ip1 = str(ip1)
    if not isinstance(ip2, str):
        ip2 = str(ip2)

    int1 = ip_to_int(ip1)
    int2 = ip_to_int(ip2)
    if int1 is None or int2 is None:
        return float("nan")

    divisor = 2.0 ** mask
    max_val = 2.0 ** (32 - mask)
    return abs(int1 / divisor - int2 / divisor) / max_val


def _haversine_km(lat1: float, lon1: float, lat2: float, lon2: float) -> float:
    """Great-circle distance (km) via the Haversine formula."""
    dlat = math.radians(lat2 - lat1)
    dlon = math.radians(lon2 - lon1)
    a = (
        math.sin(dlat / 2) ** 2
        + math.cos(math.radians(lat1))
        * math.cos(math.radians(lat2))
        * math.sin(dlon / 2) ** 2
    )
    return 6371.0 * 2.0 * math.asin(math.sqrt(a))


def _geo_distance(ip1: Optional[str], ip2: Optional[str]) -> float:
    """Normalised [0, 1] geographic distance between two IP strings.

    Uses ``analysis.geoip.lookup()`` to obtain lat/lon, then Haversine.
    Returns ``NaN`` when either IP cannot be geolocated.
    """
    if ip1 is None or ip2 is None:
        return float("nan")
    if not isinstance(ip1, str):
        ip1 = str(ip1)
    if not isinstance(ip2, str):
        ip2 = str(ip2)

    geo1 = lookup(ip1)
    geo2 = lookup(ip2)
    if geo1 is None or geo2 is None:
        return float("nan")

    km = _haversine_km(geo1["lat"], geo1["lon"], geo2["lat"], geo2["lon"])
    return km / _HALF_EARTH_CIRCUMFERENCE_KM


# ---------------------------------------------------------------------------
#  Public API
# ---------------------------------------------------------------------------

def ip_distance(
    lf: pl.LazyFrame,
    ip_columns: list[str],
    subnet_masks: Optional[list[int]] = None,
) -> pl.LazyFrame:
    """Append IP distance columns to a LazyFrame.

    For each **consecutive pair** of *ip_columns*
    (``ip_columns[0]`` vs ``ip_columns[1]``,
    ``ip_columns[2]`` vs ``ip_columns[3]``, …) computes:

    * one subnet-mask distance per entry in *subnet_masks*
    * one Haversine geographic distance

    All distances are normalised to [0, 1].  Invalid IPs or IPs not
    found in the GeoIP database produce ``NaN``.

    Parameters
    ----------
    lf:
        LazyFrame containing the IP columns.
    ip_columns:
        List of IP column names (used in consecutive pairs).
    subnet_masks:
        Subnet masks for network distance (default ``[24, 28]``).

    Returns
    -------
    pl.LazyFrame
        Input frame with ``_ip_dist_0``, ``_ip_dist_1``, …
        columns appended.

    Examples
    --------
    >>> lf = pl.LazyFrame({
    ...     "src_ip": ["8.8.8.8", "1.1.1.1"],
    ...     "dst_ip": ["8.8.4.4", "1.0.0.1"],
    ... })
    >>> result = ip_distance(lf, ["src_ip", "dst_ip"], subnet_masks=[24])
    >>> result.collect().columns
    ['src_ip', 'dst_ip', '_ip_dist_0', '_ip_dist_1']
    """
    if subnet_masks is None:
        subnet_masks = [24, 28]

    result = lf
    dist_idx = 0

    # Walk consecutive pairs
    for i in range(0, len(ip_columns) - 1, 2):
        col_a = ip_columns[i]
        col_b = ip_columns[i + 1]

        # -- subnet-mask distances --
        for mask in subnet_masks:
            name = f"_ip_dist_{dist_idx}"
            result = result.with_columns(
                pl.struct([col_a, col_b])
                .map_elements(
                    lambda s, _mask=mask, _ca=col_a, _cb=col_b: (
                        _subnet_mask_distance(s[_ca], s[_cb], _mask)
                    ),
                    return_dtype=pl.Float64,
                )
                .alias(name)
            )
            dist_idx += 1

        # -- geographic distance --
        name = f"_ip_dist_{dist_idx}"
        result = result.with_columns(
            pl.struct([col_a, col_b])
            .map_elements(
                lambda s, _ca=col_a, _cb=col_b: _geo_distance(s[_ca], s[_cb]),
                return_dtype=pl.Float64,
            )
            .alias(name)
        )
        dist_idx += 1

    return result


# ---------------------------------------------------------------------------
#  String distance
# ---------------------------------------------------------------------------


def _is_enum_like(s: pl.Series) -> bool:
    """Check if ≥50% of rows have values appearing ≥2 times."""
    total = len(s)
    if total == 0:
        return False
    non_null = s.drop_nulls()
    if len(non_null) == 0:
        return False
    vc = non_null.value_counts()
    if 'count' not in vc.columns:
        return False
    duplicated = vc.filter(pl.col('count') >= 2)
    if len(duplicated) == 0:
        return False
    duplicated_count: int = duplicated['count'].sum()  # type: ignore[union-attr]
    return duplicated_count >= 0.5 * total


def _levenshtein_normalized(s1: str | None, s2: str) -> float:
    """Normalised Levenshtein distance between *s1* and *s2*.

    Empty or ``None`` values return 1.0 (maximum distance).
    """
    if s1 is None or s1 == '':
        return 1.0
    if s2 is None or s2 == '':
        return 1.0
    s1_str = str(s1)
    s2_str = str(s2)
    max_len = max(len(s1_str), len(s2_str))
    if max_len == 0:
        return 0.0
    import Levenshtein
    return Levenshtein.distance(s1_str, s2_str) / max_len


def string_distance(lf: pl.LazyFrame, str_columns: list[str]) -> pl.LazyFrame:
    """Compute string distance columns.

    For each column in *str_columns*:

    * If **≥50%** of rows have values appearing ≥2 times, the column is
      treated as an enum: 0 when the value equals the mode, 1 otherwise.
    * Otherwise, the normalised Levenshtein distance to the mode is
      computed for each row.

    Empty or ``None`` values always produce the maximum distance (1.0).

    Parameters
    ----------
    lf:
        Input LazyFrame.
    str_columns:
        List of string column names to compute distances for.

    Returns
    -------
    pl.LazyFrame
        Input LazyFrame extended with ``_str_dist_0``, ``_str_dist_1``,
        ... columns (one per entry in *str_columns*).
    """
    df = lf.collect(streaming=True)
    dist_exprs: list[pl.Expr] = []

    for idx, col in enumerate(str_columns):
        if col not in df.columns:
            continue

        series = df[col].cast(pl.Utf8)
        enum_like = _is_enum_like(series)

        non_null = series.drop_nulls()
        if len(non_null) > 0:
            mode_val = non_null.value_counts().sort('count', descending=True).row(0)[0]
            mode_str = str(mode_val)
        else:
            mode_str = ''

        col_name = f'_str_dist_{idx}'

        if enum_like:
            dist_exprs.append(
                pl.when(pl.col(col).is_null())
                .then(pl.lit(1.0))
                .when(pl.col(col).cast(pl.Utf8) == pl.lit(mode_str))
                .then(pl.lit(0.0))
                .otherwise(pl.lit(1.0))
                .alias(col_name)
            )
        else:
            # Capture in closure for map_elements
            _mode = mode_str
            dist_exprs.append(
                pl.col(col).map_elements(
                    lambda v: _levenshtein_normalized(v, _mode),
                    return_dtype=pl.Float64,
                ).alias(col_name)
            )

    return lf.with_columns(dist_exprs)


# ---------------------------------------------------------------------------
#  Boolean distance
# ---------------------------------------------------------------------------


def bool_distance(lf: pl.LazyFrame, bool_columns: list[str]) -> pl.LazyFrame:
    """Compute boolean distance columns.

    For each column in *bool_columns*, each row gets:
    0 if its value matches the column mode, 1 otherwise.
    ``None`` values always produce 1.0 (maximum distance).

    Parameters
    ----------
    lf:
        Input LazyFrame.
    bool_columns:
        List of boolean column names to compute distances for.

    Returns
    -------
    pl.LazyFrame
        Input LazyFrame extended with ``_bool_dist_0``, ``_bool_dist_1``,
        ... columns (one per entry in *bool_columns*).
    """
    df = lf.collect(streaming=True)
    dist_exprs: list[pl.Expr] = []

    for idx, col in enumerate(bool_columns):
        if col not in df.columns:
            continue

        series = df[col]
        non_null = series.drop_nulls()
        col_name = f'_bool_dist_{idx}'

        if len(non_null) == 0:
            dist_exprs.append(pl.lit(0.0).alias(col_name))
        else:
            mode_val = non_null.value_counts().sort('count', descending=True).row(0)[0]
            dist_exprs.append(
                pl.when(pl.col(col).is_null())
                .then(pl.lit(1.0))
                .when(pl.col(col) == pl.lit(mode_val))
                .then(pl.lit(0.0))
                .otherwise(pl.lit(1.0))
                .alias(col_name)
            )

    return lf.with_columns(dist_exprs)


# ---------------------------------------------------------------------------
#  Bytes Hamming distance
# ---------------------------------------------------------------------------

_HEX_POPCOUNT: dict[str, int] = {
    '0': 0, '1': 1, '2': 1, '3': 2,
    '4': 1, '5': 2, '6': 2, '7': 3,
    '8': 1, '9': 2,
    'a': 2, 'b': 3, 'c': 2, 'd': 3, 'e': 3, 'f': 4,
    'A': 2, 'B': 3, 'C': 2, 'D': 3, 'E': 3, 'F': 4,
}


def _hex_popcount_normalized(val: str | None) -> float:
    """Return the proportion of 1-bits in a hex byte string (0.0–1.0).

    This is the Hamming distance from the all-zero bit string,
    normalised by total bits.  Returns ``NaN`` for null or
    non-hex input.
    """
    if val is None:
        return float("nan")
    s = str(val).replace(' ', '').replace('\t', '').replace('\n', '')
    if not s:
        return 0.0

    total_ones = 0
    for ch in s:
        ones = _HEX_POPCOUNT.get(ch)
        if ones is None:
            return float("nan")  # non-hex character
        total_ones += ones

    total_bits = len(s) * 4
    return total_ones / total_bits


def bytes_distance(
    lf: pl.LazyFrame,
    byte_columns: list[str],
) -> pl.LazyFrame:
    """Append normalised Hamming-distance columns for hex byte columns.

    Each byte column is converted from hex to binary, and the
    proportion of 1-bits (Hamming distance from the all-zero
    bit vector) is computed as a ``_bytes_dist_N`` Float64 column
    normalised to [0, 1].  Non-hex / null values produce ``NaN``.

    Parameters
    ----------
    lf:
        LazyFrame containing the byte columns.
    byte_columns:
        Column names whose values are hex byte strings
        (with optional whitespace separators, e.g. ``"A3 F2"``).

    Returns
    -------
    pl.LazyFrame
        Input frame with ``_bytes_dist_0``, ``_bytes_dist_1``, …
        appended.  Columns not present in *lf* are silently skipped.

    Examples
    --------
    >>> df = pl.DataFrame({"payload": ["A3F2", "FFFF", None]})
    >>> result = bytes_distance(df.lazy(), ["payload"])
    >>> result.collect()
    shape: (3, 2)
    ┌─────────┬────────────────┐
    │ payload ┆ _bytes_dist_0  │
    │ ---     ┆ ---            │
    │ str     ┆ f64            │
    ╞═════════╪════════════════╡
    │ A3F2    ┆ 0.625          │
    │ FFFF    ┆ 1.0            │
    │ null    ┆ NaN            │
    └─────────┴────────────────┘
    """
    existing = set(lf.collect_schema().names())
    result = lf

    for i, col_name in enumerate(byte_columns):
        if col_name not in existing:
            continue
        result = result.with_columns(
            pl.col(col_name)
            .cast(pl.Utf8)
            .map_elements(_hex_popcount_normalized, return_dtype=pl.Float64)
            .alias(f'_bytes_dist_{i}')
        )

    return result


# ---------------------------------------------------------------------------
#  Timestamp FFT phase distance
# ---------------------------------------------------------------------------

import numpy as np

# ── threshold: minimum fraction of repeated intervals to use FFT path ──
_FFT_PERIODICITY_THRESHOLD = 0.05
# ── minimum samples before FFT is attempted ──
_FFT_MIN_SAMPLES = 4


def _parse_timestamp_to_epoch(series: pl.Series) -> pl.Series:
    """Convert a timestamp series to Unix epoch seconds (float64).

    Handles int/float (assumed already epoch), pl.Datetime/pl.Date
    (ns→seconds), and str/pl.Utf8 (strptime→ns→seconds, with direct
    float-cast fallback on parse failure).
    """
    dtype = series.dtype
    if dtype in (pl.Utf8, pl.String):
        try:
            parsed = series.str.strptime(pl.Datetime, format=None, strict=False)
        except Exception:
            return series.cast(pl.Float64)
        return parsed.cast(pl.Int64).cast(pl.Float64) / 1_000_000_000.0
    elif dtype in (pl.Datetime, pl.Date):
        return series.cast(pl.Int64).cast(pl.Float64) / 1_000_000_000.0
    else:
        return series.cast(pl.Float64)


def _compute_fft_distances(values: np.ndarray) -> np.ndarray:
    """Compute per-row phase distances for a 1-D array.

    Sorts *values*, checks for periodicity, then either:
    * FFT path: ``np.fft.fft()`` → ``np.angle()``, mapped back to
      original (unsorted) row order via rank lookup.
    * Statistical fallback: interval z-scores mapped to row order.

    Returns an ``(n,)`` float64 array, normalised to zero-mean
    unit-variance.
    """
    n = len(values)
    if n < _FFT_MIN_SAMPLES:
        return np.full(n, np.nan, dtype=np.float64)

    # Sort once; record rank→position mapping for unsorting later
    sorted_order = np.argsort(values)
    ranks = np.empty(n, dtype=np.int64)
    ranks[sorted_order] = np.arange(n)
    sorted_vals = values[sorted_order]

    # ── periodicity check ──────────────────────────────────────────
    intervals = np.diff(sorted_vals)
    unique_int, counts = np.unique(
        np.round(intervals, 6), return_counts=True
    )
    repeat_ratio = float(np.sum(counts[counts > 1])) / max(len(intervals), 1)

    if repeat_ratio >= _FFT_PERIODICITY_THRESHOLD:
        # ── FFT phase path ─────────────────────────────────────────
        centered = sorted_vals - np.mean(sorted_vals)
        fft_result = np.fft.fft(centered)
        phases = np.angle(fft_result)
        dist_values = phases[ranks]
    else:
        # ── statistical fallback ───────────────────────────────────
        mean_int = float(np.mean(intervals))
        std_int = float(np.std(intervals))
        if std_int > 1e-10:
            z_scores = (intervals - mean_int) / std_int
        else:
            z_scores = np.zeros_like(intervals)

        # Row at sorted position r gets the z-score of interval r-1
        # (preceding interval); row at position 0 reuses last interval.
        dist_values = np.empty(n, dtype=np.float64)
        for r in range(n):
            if r > 0:
                dist_values[ranks[r]] = z_scores[r - 1]
            else:
                dist_values[ranks[0]] = (
                    z_scores[-1] if n > 1 else 0.0
                )

    # Normalise to zero-mean unit-variance
    d_std = float(np.std(dist_values))
    d_mean = float(np.mean(dist_values))
    if d_std > 1e-10:
        dist_values = (dist_values - d_mean) / d_std
    else:
        dist_values = dist_values - d_mean

    return dist_values


def timestamp_fft_distance(
    lf: pl.LazyFrame, ts_columns: list[str]
) -> pl.LazyFrame:
    """Append FFT-based phase-distance columns for timestamp columns.

    For each column in *ts_columns* that exists in the schema:

    1. Parse values to Unix epoch float seconds (handles int, float,
       datetime strings, and ``pl.Datetime``).
    2. Collect the column to numpy, sort, then run ``np.fft.fft``
       followed by ``np.angle`` to extract per-row phase values.
    3. If fewer than 5 % of timestamp intervals repeat (weak
       periodicity), fall back to statistical features derived
       from interval mean, standard deviation, and z-scores.

    Columns are appended as ``_ts_dist_0``, ``_ts_dist_1``, … in the
    order *ts_columns* are supplied.  Rows whose timestamp is null
    receive ``NaN``.

    Parameters
    ----------
    lf : pl.LazyFrame
        The input LazyFrame with at least one timestamp column.
    ts_columns : list[str]
        Column names to process.

    Returns
    -------
    pl.LazyFrame
        A new LazyFrame with ``_ts_dist_N`` columns appended.

    Examples
    --------
    >>> import polars as pl
    >>> from analysis.distance import timestamp_fft_distance
    >>> df = pl.DataFrame({
    ...     'ts': [0, 60, 120, 180, 240, 300],
    ... })
    >>> result = timestamp_fft_distance(df.lazy(), ['ts'])
    >>> result.collect().columns
    ['ts', '_ts_dist_0']
    """
    schema = lf.collect_schema()
    available = [c for c in ts_columns if c in schema.names()]
    if not available:
        return lf

    # Collect all needed columns once
    df = lf.select(available).collect(streaming=True)
    n_rows = len(df)

    dist_arrays: list[np.ndarray] = []

    for col in available:
        series = df[col]
        ts_numeric = _parse_timestamp_to_epoch(series)

        # Identify valid (non-null) rows
        valid = ~ts_numeric.is_null()
        valid_indices = np.where(valid.to_numpy())[0]
        values = ts_numeric.filter(valid).to_numpy().astype(np.float64)

        if len(values) < _FFT_MIN_SAMPLES:
            continue

        phase_values = _compute_fft_distances(values)

        # Build full-length array; invalid rows stay NaN
        full_dist = np.full(n_rows, np.nan, dtype=np.float64)
        full_dist[valid_indices] = phase_values
        dist_arrays.append(full_dist)

    if not dist_arrays:
        return lf

    # Re-materialise the original frame and bolt on distance columns
    result_df = lf.collect(streaming=True)
    for i, arr in enumerate(dist_arrays):
        result_df = result_df.with_columns(
            pl.Series(f'_ts_dist_{i}', arr)
        )

    return result_df.lazy()


# ---------------------------------------------------------------------------
#  Normalization — StandardScaler (step 2 of preprocessing pipeline)
# ---------------------------------------------------------------------------

from sklearn.preprocessing import StandardScaler  # noqa: E402


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()


# ---------------------------------------------------------------------------
#  SVD noise removal (step 3 of preprocessing pipeline)
# ---------------------------------------------------------------------------

from sklearn.decomposition import TruncatedSVD  # noqa: E402


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


# ---------------------------------------------------------------------------
#  Master preprocessing pipeline
# ---------------------------------------------------------------------------


def run_preprocessing_pipeline(
    lf: pl.LazyFrame,
    column_types: dict | None = None,
) -> pl.LazyFrame:
    """Full preprocessing pipeline for UMAP / clustering.

    Steps (executed sequentially):
    1. **Column type detection** — uses
       :func:`analysis.type_classifier.classify_schema` unless
       *column_types* is explicitly provided.
    2. **Distance computation** — IP distance, string distance,
       bool distance, bytes distance, and timestamp FFT distance
       computed per-column-type and appended as ``_ip_dist_*``,
       ``_str_dist_*``, ``_bool_dist_*``, ``_bytes_dist_*``,
       ``_ts_dist_*`` columns.
    3. **Normalisation** — :func:`normalize_features` (StandardScaler
       → μ=0, σ=1).
    4. **SVD denoise** — :func:`svd_denoise` (TruncatedSVD keeping
       95 % variance).

    Parameters
    ----------
    lf:
        Input LazyFrame with raw data columns.
    column_types:
        Optional ``{col_name: DataType}`` mapping.  When ``None``,
        :func:`analysis.type_classifier.classify_schema` is called
        to auto-detect types.

    Returns
    -------
    pl.LazyFrame
        LazyFrame with only the SVD-denoised feature columns
        (``_svd__norm__*_dist_*`` etc.) plus original row anchors.
    """
    if column_types is None:
        from analysis.types import classify_schema
        column_types = classify_schema(lf)

    ip_cols: list[str] = []
    str_cols: list[str] = []
    bool_cols: list[str] = []
    byte_cols: list[str] = []
    ts_cols: list[str] = []
    float_cols: list[str] = []
    int_cols: list[str] = []

    for col_name, ct in column_types.items():
        if ct.name == "IPv4":
            ip_cols.append(col_name)
        elif ct.name == "字符串":
            str_cols.append(col_name)
        elif ct.name == "布尔":
            bool_cols.append(col_name)
        elif ct.name == "字节":
            byte_cols.append(col_name)
        elif ct.name == "时间戳":
            ts_cols.append(col_name)
        elif ct.name in ("浮点", "整数"):
            # Float/Int are kept as direct numeric features
            if ct.name == "浮点":
                float_cols.append(col_name)
            else:
                int_cols.append(col_name)

    result = lf

    # -- Step 2: Distance computation --
    if ip_cols and len(ip_cols) >= 2:
        result = ip_distance(result, ip_cols)
    if str_cols:
        result = string_distance(result, str_cols)
    if bool_cols:
        result = bool_distance(result, bool_cols)
    if byte_cols:
        result = bytes_distance(result, byte_cols)
    if ts_cols:
        result = timestamp_fft_distance(result, ts_cols)

    # Collect numeric feature columns present after distance computation
    schema = result.collect_schema()
    all_dist_cols = [
        c for c in schema.names()
        if (c.startswith('_ip_dist_') or c.startswith('_str_dist_')
            or c.startswith('_bool_dist_') or c.startswith('_bytes_dist_')
            or c.startswith('_ts_dist_'))
    ]
    # Also include raw FLOAT/INT columns as direct features
    numeric_raw = [
        c for c in (float_cols + int_cols)
        if c in schema.names()
    ]

    feature_cols = all_dist_cols + numeric_raw
    if not feature_cols:
        return result

    # -- Step 3: Normalisation --
    result = normalize_features(result, feature_cols)

    # -- Step 4: SVD denoise (REPLACES noisy features with cleaned ones) --
    norm_cols = [f'_norm_{c}' for c in feature_cols]
    result, noise_profile = svd_denoise(result, norm_cols)

    return result, noise_profile


# ---------------------------------------------------------------------------
#  UMAP dimensionality reduction (3D → 2D fallback)
# ---------------------------------------------------------------------------


def umap_reduce(
    lf: pl.LazyFrame,
    feature_cols: list[str],
    n_components: int = 3,
) -> tuple[pl.LazyFrame, dict]:
    """UMAP dimensionality reduction with automatic 3D→2D fallback.

    Trains on a ≤10K-row sample (when *lf* exceeds that limit), then
    batch-transforms the remaining rows in 1K-row batches.

    If ``n_components=3`` fails (e.g. too few rows or numerical issues),
    falls back to ``n_components=2`` automatically.

    Parameters
    ----------
    lf:
        LazyFrame containing the numeric feature columns.
    feature_cols:
        Numeric column names to use for UMAP projection.  Columns not
        present in *lf* are silently skipped.
    n_components:
        Target dimensionality (2 or 3).  Default 3 with fallback to 2.

    Returns
    -------
    tuple[pl.LazyFrame, dict]
        - LazyFrame: input data extended with ``_umap_x``, ``_umap_y``
          (and ``_umap_z`` when ``n_components ≥ 3``).
        - dict: metadata with keys ``n_components`` and ``reducer``.
    """
    import numpy as np
    from sklearn.preprocessing import StandardScaler
    import umap
    import logging

    logger = logging.getLogger(__name__)
    MAX_TRAIN = 10000
    BATCH_SIZE = 1000

    available = [c for c in feature_cols if c in lf.collect_schema().names()]
    if not available:
        return lf, {'n_components': 0, 'reducer': 'umap', 'error': 'no feature columns available'}

    df = lf.collect(streaming=True)
    n_total = len(df)

    if n_total < 2:
        return lf, {'n_components': 0, 'reducer': 'umap', 'error': 'too few rows'}

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

    scaler = StandardScaler()
    if n_total > MAX_TRAIN:
        rng = np.random.RandomState(42)
        idx_sample = rng.choice(n_total, size=MAX_TRAIN, replace=False)
        mat_sample = mat[idx_sample]
        scaler.fit(mat_sample)
        mat_sample_scaled = np.nan_to_num(scaler.transform(mat_sample), nan=0.0)
    else:
        scaler.fit(mat)
        mat_sample_scaled = np.nan_to_num(scaler.transform(mat), nan=0.0)

    # ── Train UMAP; fallback 3→2 on failure ──
    reducer = None
    coords = None
    effective_n = n_components

    # Try requested n_components first, then fall back to 2 (no duplicates)
    attempts = sorted(set((n_components, 2)), reverse=True)
    for attempt_n in attempts:
        if attempt_n != n_components:
            logger.warning(f'UMAP {n_components}D failed, falling back to 2D')
        effective_n = attempt_n
        try:
            reducer = umap.UMAP(
                n_components=effective_n, random_state=42,
                n_neighbors=15, min_dist=0.1, metric='euclidean',
            )
            if n_total > MAX_TRAIN:
                reducer.fit(mat_sample_scaled)
            else:
                coords = reducer.fit_transform(mat_sample_scaled)
            break
        except Exception as e:
            if attempt_n == 2:
                raise
            logger.warning(f'UMAP {attempt_n}D failed: {e}')

    # ── Batch transform if dataset was larger than training sample ──
    if n_total > MAX_TRAIN and reducer is not None:
        coords_list: list[np.ndarray] = []
        for start in range(0, n_total, BATCH_SIZE):
            end = min(start + BATCH_SIZE, n_total)
            mat_batch = scaler.transform(mat[start:end])
            mat_batch = np.nan_to_num(mat_batch, nan=0.0, posinf=0.0, neginf=0.0)
            coords_list.append(reducer.transform(mat_batch))
        coords = np.vstack(coords_list)

    if coords is None:
        return lf, {'n_components': 0, 'reducer': 'umap', 'error': 'UMAP produced no output'}

    # ── Build result LazyFrame with coordinate columns ──
    coord_series = [
        pl.Series('_umap_x', coords[:, 0]),
        pl.Series('_umap_y', coords[:, 1]),
    ]
    if effective_n >= 3 and coords.shape[1] >= 3:
        coord_series.append(pl.Series('_umap_z', coords[:, 2]))

    df_result = df.with_columns(coord_series)

    metadata: dict = {
        'n_components': effective_n,
        'reducer': 'umap',
    }

    return df_result.lazy(), metadata


# ---------------------------------------------------------------------------
#  Per-cluster SVD feature extraction
# ---------------------------------------------------------------------------


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.
    """
    import numpy as np
    from sklearn.decomposition import TruncatedSVD

    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