"""Saturación realista de nubes brillantes."""

import numpy as np
from scipy.ndimage import binary_dilation, generate_binary_structure

from ..config import SaturationConfig


def apply_cloud_saturation(
    image: np.ndarray,
    labels: np.ndarray | None,
    config: "SaturationConfig",
) -> np.ndarray:
    """
    Simula saturación de sensor en nubes brillantes.

    En MSS real, píxeles muy brillantes saturan el detector,
    causando que múltiples bandas tengan valores idénticos o muy similares.

    Args:
        image: (C, H, W) float32
        labels: (H, W) uint8 con clases (2=cloud, 1=thin)
        config: SaturationConfig

    Returns:
        Imagen con saturación aplicada
    """
    if not config.enabled or labels is None:
        return image

    C, H, W = image.shape
    result = image.copy()

    # Detectar nubes brillantes (clase 2 = cloud, o clase 1 si merge_cloud_classes)
    cloud_mask = (labels == 2) | (labels == 1)

    if not cloud_mask.any():
        return result

    # Detectar píxeles brillantes dentro de nubes
    brightness = np.mean(image[:3], axis=0)  # Promedio de primeras 3 bandas
    bright_threshold = config.brightness_threshold.sample()
    bright_clouds = cloud_mask & (brightness > bright_threshold)

    if not bright_clouds.any():
        return result

    # Expandir regiones de saturación (simula bleeding del sensor)
    struct = generate_binary_structure(2, 2)
    dilate_iterations = config.dilate_pixels.sample()
    saturated_mask = binary_dilation(bright_clouds, struct, iterations=dilate_iterations)
    saturated_mask = saturated_mask & cloud_mask  # Solo dentro de nubes

    if not saturated_mask.any():
        return result

    # Aplicar saturación con patrón realista
    sat_value = config.saturation_value.sample()
    band_variation = config.band_variation.sample()

    # Generar valores de saturación por banda (similares pero no idénticos)
    sat_values = np.array(
        [sat_value + np.random.uniform(-band_variation, band_variation) for _ in range(C)]
    )

    # Las primeras bandas tienden a saturar más uniformemente
    if config.uniform_first_bands:
        # B1, B2, B3 comparten valor base (patrón observado en MSS real)
        base_sat = sat_values[0]
        sat_values[0] = base_sat
        sat_values[1] = base_sat + np.random.uniform(-0.01, 0.01)
        sat_values[2] = base_sat + np.random.uniform(-0.02, 0.02)
        # B4 (NIR) puede variar más
        sat_values[3] = sat_value + np.random.uniform(-band_variation, band_variation)

    # Aplicar saturación gradual (no abrupta)
    for c in range(C):
        # Blend suave hacia valor de saturación
        blend_factor = config.blend_strength.sample()
        result[c, saturated_mask] = (
            result[c, saturated_mask] * (1 - blend_factor) + sat_values[c] * blend_factor
        )

    # Añadir ruido mínimo para evitar valores exactamente iguales
    noise_std = config.noise_std.sample()
    if noise_std > 0:
        noise = np.random.normal(0, noise_std, result[:, saturated_mask].shape)
        result[:, saturated_mask] += noise

    return np.clip(result, 0, 1).astype(np.float32)
