from __future__ import annotations

import numpy as np


def _band_score(band: np.ndarray) -> float:
    band = band.astype("float32")
    finite = band[np.isfinite(band)]
    if finite.size == 0:
        return 0.0
    variance = float(np.var(finite))
    gradient = np.mean(np.abs(np.gradient(finite))) if finite.size > 1 else 0.0
    return variance + gradient


def select_features(bands: np.ndarray, max_features: int = 6) -> list[int]:
    if bands.ndim != 3:
        raise ValueError("Bands array must have shape (bands, height, width).")

    scores = [
        _band_score(bands[idx]) for idx in range(bands.shape[0])
    ]
    order = np.argsort(scores)[::-1]

    selected: list[int] = []
    for idx in order:
        if len(selected) >= max_features:
            break
        if not selected:
            selected.append(int(idx))
            continue
        candidate = bands[int(idx)].reshape(-1)
        correlations = []
        for sel in selected:
            ref = bands[sel].reshape(-1)
            if np.std(ref) == 0 or np.std(candidate) == 0:
                corr = 1.0
            else:
                corr = float(np.corrcoef(ref, candidate)[0, 1])
            correlations.append(abs(corr))
        if correlations and max(correlations) > 0.92:
            continue
        selected.append(int(idx))

    if not selected:
        selected = [0]

    return selected
