"""UMAP dimensionality reduction (3D → 2D fallback)."""

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


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``.
    """
    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