from __future__ import annotations

import numpy as np


def percentile_stretch(band: np.ndarray, low: float = 2, high: float = 98) -> np.ndarray:
    if band.size == 0:
        return band
    band = band.astype("float32")
    valid = band[np.isfinite(band)]
    if valid.size == 0:
        return band
    lo = np.percentile(valid, low)
    hi = np.percentile(valid, high)
    if hi <= lo:
        hi = lo + 1.0
    stretched = (band - lo) / (hi - lo)
    return np.clip(stretched, 0, 1)


def to_uint8(band: np.ndarray, low: float = 2, high: float = 98) -> np.ndarray:
    stretched = percentile_stretch(band, low, high)
    return (stretched * 255).astype("uint8")


def prepare_rgb(
    bands: np.ndarray,
    indices: tuple[int, int, int],
    stretch: tuple[float, float] = (2, 98),
) -> np.ndarray:
    r = to_uint8(bands[indices[0]], *stretch)
    g = to_uint8(bands[indices[1]], *stretch)
    b = to_uint8(bands[indices[2]], *stretch)
    return np.stack([r, g, b], axis=-1)


def prepare_gray(bands: np.ndarray, index: int, stretch: tuple[float, float] = (2, 98)) -> np.ndarray:
    return to_uint8(bands[index], *stretch)


def flatten_features(bands: np.ndarray, mask: np.ndarray | None = None) -> tuple[np.ndarray, np.ndarray]:
    bands = bands.astype("float32")
    features = bands.reshape(bands.shape[0], -1).T  # (pixels, bands)
    if mask is None:
        mask = np.isfinite(features).all(axis=1)
    else:
        mask = mask.reshape(-1) & np.isfinite(features).all(axis=1)
    return features[mask], mask


def restore_map(
    values: np.ndarray, mask: np.ndarray, shape: tuple[int, int], fill_value: float | int = np.nan
) -> np.ndarray:
    full = np.full(mask.shape[0], fill_value, dtype=values.dtype)
    full[mask] = values
    return full.reshape(shape)
